UltraTWD: Optimizing Ultrametric Trees for Tree-Wasserstein Distance
Authors: Fangchen Yu
Yanzhen Chen
Jiaxing Wei
Jianfeng Mao
Wenye Li
Qiang Sun
Venue: ICML 2025 poster (2025)
Keywords: 1-Wasserstein Distance
Tree-Wasserstein Distance
Ultrametric Tree
Optimization
Primary area: general_machine_learning
TL;DR: We propose UltraTWD, a novel framework for accurate tree-Wasserstein distance computation using optimized ultrametric trees, significantly improving performance in text-based applications.

Abstract: The Wasserstein distance is a widely used metric for measuring differences between distributions, but its super-cubic time complexity introduces substantial computational burdens. To mitigate this, the tree-Wasserstein distance (TWD) offers a linear-time approximation by leveraging a tree structure; however, existing TWD methods often compromise accuracy due to suboptimal tree structures and edge weights. To address it, we introduce UltraTWD, a novel unsupervised framework that simultaneously optimizes both ultrametric tree structures and edge weights to more faithfully approximate the cost matrix. Specifically, we develop algorithms based on minimum spanning trees, iterative projection, and gradient descent to efficiently learn high-quality ultrametric trees. Empirical results across document retrieval, ranking, and classification tasks demonstrate that UltraTWD achieves superior approximation accuracy and competitive downstream performance. Code is available at: https://github.com/NeXAIS/UltraTWD.

UltraTWD: Optimizing Ultrametric Trees for Tree- Wasserstein Distance
Fangchen Yu Yanzhen Chen Jiaxing Wei!
Jianfeng Mao
Wenye Li Qiang Sun 25*
Abstract
Euclidean space
Tree space
The Wasserstein distance is widely used metric for measuring differences between distributions, but its super-cubic time complexity introduces substantial computational burdens. To mitigate this, the tree- Wasserstein distance (TWD) offers linear-time approximation by leveraging a tree structure; however, existing TWD methods of- ten compromise accuracy due suboptimal tree structures and edge weights. To address it, we in- troduce UltraTWD, a novel unsupervised frame- work that simultaneously optimizes both ultra- metric tree structures and edge weights to more faithfully approximate the cost matrix. Specifi- cally, we develop algorithms based on minimum spanning trees, iterative projection, and gradient descent to efficiently learn high-quality ultramet- ric trees. Empirical results across document re- trieval, ranking. and classification tasks demon- strate that UlD achieves superior approx- imation accuracy and competitive downstream performance. Code is available at: https: httpehaw..cco//eee////tttteee
d(x1, x4)
X
X
(x,,,,
X3
X2
X3
(a) Cost d(x,, xy) for Wasser- stein distance
(b) Tree distance dy(x,, xy) for tree- Seee-assersttin distance
Figure I. The tree-Wer Wasserstein distance embeds data points into tree and replaces the cost d(æ;;, mzy) with tree distance (1r,, mz))
Several approximation methods have been proposed to mit- igate the computational challenges of the Wasserstein dis- tance. The Sinkhorn distance (Cuturi, 2013) reduces the complexity to O(n2) by introducing an entropic regulariza- tion term. The Bicee-Wasssesstenn distance (SWD) (Rabin et al., 2012) achieves Onn log complexity by projecting nigh-dimensional data onto one-dimensional spaces and computing distances over these ID projections. The tree- Wasserstein distance (TWD) (Le et al., 2019) further im- proves efficiency by learning a rooted tree, enabling (n)) complexity, as illustrated in Fig. This work focuses on advancing the TWD methodology.
1. Introduction
Existing TWD methods can be categorized into four cate- gories, each with notable limitations. (1) Tree-construction methods: QuadTree (Indyk & Thaper, 2003) and Clus- terTree (Le et al., 2019) construct trees by partitioning hyper- cubes through recursive division or farthest-point clustering (Gonzalez, 1985). However, their edge weights, determined by the depth of nodes, often suboptimal. (2) Weight- optimized methods: qTWD and cTWD (Yamada al., 2022) refine the edge weights of QuadTreee or ClusterTree using Lasso regression but fail to optimize the tree structures themselves, resulting in performance heavily dependent on the randomness of the initial tree. (3) Tree-sliced meth- ods: Sliced-gTWD and cTWD (Yamada et al., 2022; Otao & Yamada, 2023) improve stability by averaging multiple random trees, yet they cannot optimize tree structures. (4) Supervised method: UltraTree (Chen et al., 2024) learns an ultrametric tree by regressing TWD against the true Wasser- stein distance. However, this approach requires extensive training data with precomputed Wasserstein distances as labels, which makes computationally expensive.
The Wasserstein distance (Villani, 2008) is a powerful met- ric for measuring dissimilarities between probability distri- butions, with various applications (Takezawa et al., 2021; Khrulkov et al., 2023) A well-known example, the Word Mover's Distance (Kusner et al., 2015), computes distances between documents modeled as bag-of-words distributions over n word embeddings. Despite its effectiveness, the Wasserstein distance suffers from a high computational com- plexity of On 3 log n), limiting its practical applicability.
'The Chinese University Hong Kong, Shenzhen MMohamed bin Zayed University of Artificial Intelligence 'Shenzhen Research Institute of Big Data "The Hong Kong University of Science and Technology (Guangzhou) University of Toronto. Correspondence to: Wenye Li <wenyeli hwwst-gz..ccooc wwwssttzz.euu..n... Qiang Sun <qsun- stats stat@@gmail.com comm.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

UWD: Optimizing Ultrametric Trees for Tree-Wasserstein Distance
Given these challenges, unsupervised methods fail opti- mize tree structures and are susceptible to random initializa- tion, while the supervised method depends expensive la- beled training data. These limitations motivate us to develop a new framework that optimizes both tree structure and edge weights without requiring training data. To achieve this, we propose UltraTWD, a novel unsupervised framework that simultaneously learns the ultrametric tree structure and edge weights, delivering enhanced accuracy. To ensure the tree distance dr closely approximates the cost d, our method constructs an ultrametric tree equipped with an ultrametric Dr = [dr( 2y)] E _ that aligns with the cost ma- trix D = [d( x2y] E Rn*x. Specifically, this is achieved by solving ultrametric nearness problems, such as
Definition 1 (I-Wasserstein Distance). Given two dis-
tributions is Ei=1 ardms and =E dz0ng the I- Wasserstein distance between and is defined as:
n
Wi v) = min GEL(UM))
(1)
1J=1
where the cost d( 22²) represents the distance between Dy and zjj, and y is the transport plan in I(u, v) defined as
I(H, v) = 1- E RAXa
En = =0;, Σ Yis = =bg}, (2) "is
i
where a =(a,,-...
[a1,.... (un) and b = 111,.... bn) are mass
vectors satisfying the total mass E = Σド= 1.
Eq. (1) minimizes the total cost of transforming distribution
i into v, solvable using linear programming algorithm with computational complexity of On log (Villani, 2008). With Euclidean distance as the cost, W1 corresponds to the Word Mover's Distance (Kusner et al., 2015).
min Dr Dr
D1 or min Dr - DIE Dr
Throughout the optimization, both the tree structure and edge weights are iteratively refined, resulting in more accurate and robust tree- rre-Wasserstein distance.
2.2. Tree- Wasserstein Distance
Our contributions are summarized as follows:
The tree- rreeeasserstein distance (TWD) efficiently approx- imates the I-Wasserstein distance using tree structure. Given the vocabulary X = [z1,
We introduce UtraTWWD, new framework for tree- Wasserstein distance that leverages the ultrametric property. To our knowledge, this is the first unsupervised framework to simultaneously optimize tree structure and edge weights, overcoming the reliance of prior work on fixed tree struc- tures, multiple trees, or extensive training data.
En, E dx each word vector x; is embedded into a corresponding leaf node I in the tree (Fig. 2). This tree structure enables TWD to provide closed-form solution and achieve efficient com- putation with linear complexity of O(n), where n is both the vocabulary size and number of leaves.
We formulate ultrametric nearness problems to construct ultrametric trees that closely approximate a given cost ma- trix. To solve these problems efficiently, we propose three algorithms: (1) a fast minimum spanning tree-based ap- proach, (2) an iterative projection method, and (3) gradient baseent-based method, all yielding high-quality solutions.
Distribution μ
Tree T
Te M(Te) = a1 + d2
a1
X1
02
X2
We
Ue
Our UtraTWD achieves the lowest estimation errors among unsupervised methods and even surpasses the super- vised method across four benchmark datasets. Furthermore, our framework demonstrates strong empirical performance in document retrieval, ranking, and classification tasks, high- lighting its practicality for Wasserstein-based applications.
a mass
Xn
Xj: vector E Rdd
4: leaf node
Figure 2. Illustration of the tree- see-Wasserstein distance.
Distance). Given a tree T with tree distance dy, the tree- tree- Wasserstein distance between distributions μ and v is defined as ( (Yamada el al., 2022);
Definition 2 TTee-Wassersse
2. Preliminaries
Wr(H,v)
min Σ (1d SET(IUM) 5,j=1
(3)
2.1. 1- I-Wasersttiin Distance
The Wasserstein distance quantifies the difference between distributions. In text analysis, each word is represented by an embedding vector æ; RRd, with the vocabulary X- [x1,.... 10i,....28 _ E Bdxn containing n words. A document is modeled as a normalized bag-of-words distribution, expressed as Σi= a.da,, where òg, is the Dirac delta at Ti. mass a; represents the frequency of the i-th word in μu, and Σi= a; 1. In this paper, the Wasserstein distance refers to the W-Wasserstein distance, defined as follows:
where dr( æzj) represents the shortest path distance be- tween leaves 1; and Lú, and ry is the transport plan as defined in Eq. (2). The TWD has the following analytical form: form:
W (1, v) =Ew, w. -I((TT)) ((Te)l eET
(4)
where We denotes the weight of edge E T, and ((Te), (T.) are the total masses within the subtree Ta rooted at the deeper node of edge as illustrated in Fig. 2.
2

---

UWt Optimizing Ultrametric Trees for Tree-Wasserstell Distance
Input
UltraWWD-MST
minp, 10
Dl
Output
Distrihution Pi
MST(D)
MST(D
Dj
constant
UtrraWD-IIP
mingy Dr
DI
Distribution v v22,
b,
X1
...
movement
y
projection
D
MS
Xn Ultrametric Tree T
b-
X2
minp,
|Dy DI
update Hy,
update Dy
update
Cost Matrix
D
Tree Distance Matrix
D = [d(x, x,)]
Dr (x,
Figure diagram the proposed UltaTWD framework tree- Wasserstein distance computation. UltaTWD-MSS applies minimum spanning trees (MST) to minimize Dr Dl with an optimal guarantee. To minimize Dr - DII UltrTWW--IP employs iterative projection (IP) on the distance matrix, while UtrATWD-GD uses gradient descent (GD) for efficient tree structure optimization.
3. Unsupervised Tree- Wasserstein Distance Based on Ultrametric Trees
The fundamental mismatch between Dy and D creates gap between WT and Wi. Bridging this gap requires finding tree metric Dr that closely approximates the given cost matrix D, which serves as the motivation for our framework.
In this section, we first analyze the approximation gap be- tween Wr and W1, motivating our formulation of two ultra- metric nearness problems. We then present the UltraTWD framework with three unsupervised approaches for ultramet- ric optimization, concluding with the algorithm analysis.
3.2. Formulating Ultrametric Nearness Problems
To closely approximate D, we first formulate the tree- metric nearness problem, which aims to find the nearest tree metric Dy to the given cost matrix E R=x"
3.1. Bridging the Gap between WT and W1
The discrepancy between the tree- Wasserstein distance Wr and the 1-Wasserstein distance W1 primarily stems from the difference between the tree distance dr and the cost d. As defined in Eqs. (1) and (3), if dr (zz xz)) d(æ;, 2bj) for all pairs (æ4, my)). then (r(,, = W1 WI( v) for any distributions and sale//yYn (Yamada et al., 2022). However, this ideal scenario is rarely achievable because the tree distance matrix Df = [dr (xa, æ3)] ] E = must be a tree metric, whereas the cost matrix D ((x,,,,)) E RRxn generally does NOT conform to tree metric.
min ER
| D(I,
dyf + dhs S max dg. Ce
4 dfp)). d- 0, di = dy. 2 0,00 S si
s.t.
where di represents dr xz). However, solving this problem is intractable due to its O constraints involving four points and the non-convex nature of the optimization.
To simplify the problem, focus ultrametrics, special subset of tree metrics. In rooted tree T, the tree distance dr (x4, 2b) between two leaves 4 and ly is defined as the height of their least common ancestor (LCA):
Definition 3 (Distance and Tree Metric (Semple & Steel, 2003)). Consider symmetric, non-negative matrix D e RaxR with zero diagonal values. For all 1 S j,j,k, S n. D is a distance metric if it satisfies the triangle inequality:
dr( my) = h(LCA(1,;)) 1)).
(8)
dig S dik + dypi
(5)
E Definition 4 (Ultrametric (Semple & Steel, 2003)). tree metric Dr Raxa is an ultrametric if and only it satis- fies the strong triangle inequality:
This definition induces an ultrametric Dr = [dr(2;; æ.j)] Ruxn and is referred to as an ultrametric tree.
a tree metric it satisfies the four-point condition:
di + du S max + dit, da + d.
(6)
From Definition as setting / k in Eq. (6) implies Eq. (5). However,
tree metric must be distance metric,
drf < max(d d5eh. vigijjksn.
(9)
distance metric is not necessarily tree metrie. For example,
0
3
5
47
For ultrametric tree, the edge weight is defined as w. 1/ (h(u,) ) - h(ve)). where ue (ve) is the parent (child) node; see Fig. 2. Then, d7(2;, 25)) in Eq. (8) equals the shortest path distance between 4, and 1y. consistent with Definition 2.
0
4
the Euclidean distance matrix
satisfies
4
0
3
4
5
3
0_
all triangle incqualities but fails the four Courtpoint condition,
due to d13 + dz2= = 10 > max (d12 + d34, d1s + d23).
See Appendix A.3 for ultrametric background.
3

---

UWt: Optimizing Ultrametric Trees for Tree-Wasserstein Distance
Focusing on ultrametrics simplifies Problem (7) into the following ultrametric nearness problems, which aim to find the nearest ultrametric Dr to given cost matrix D. Specifically, we consider two optimization objectives:
3.4. Optimizing Ultrametries with Frobenius Norm: Iterative Projection Method
Unlike the L -nearest ultrametric, which has closed-form solution, the Frobenius- norm-based ultrametric nearness problem (Eq. (11)) is NP-hard (Křivánek, 1988). Prior meth- ods either relax the strong triangle inequalities into uncon- (11) strained objectives with cluster-based regularization (Chier- chia & Perret, 2019; Cohen-Addad et al., 2018; Chatziafratis et al., 2018), limiting their use to hierarchical clustering. or focus on alternative triplet relations like LCA relations (Wang & Wang, 2020; Lin et al., 2025), which do not guar- antee the ultrametric property. In contrast, we directly en- force the triplet constraint d' S max dfke} and solve Eq. (11) through a matrix optimization method.
min Dy EBaxn
|P - D1
(10)
or,
min E
Dr DII?;
subject to the constraints:
sj S max{ (dee d -f = 22 Vi,j, E [n].
The infinity norm is defined as || == m [X,] and minimizes the largest entrywise deviation between Dr and D. By controlling the maximum error, this objective ensures that Dy avoids significant outliers.
The Frobenius norm is defined as ||X := Eig X and minimizes the overall squared error between Dr and D. This objective distributes the error more evenly across all entries, resulting in good average approximation.
Although exact minimization extremely challenging, we propose an approximate solution using iterative projections (IP). The key idea is to iteratively enforce the ultrametric constraints for each triplet (i.j, k).
As illustrated in Fig. 3, we propose the UltraWWD frame- work to minimize these two objectives, introducing three approaches for optimizing ultrametrics, detailed as follows.
First, we solve the sub-problem the subset Payk imposes constraints for a single triplet:
min Dr where Ef
Pays == (Dr E S" | di S max dg.),,
3.3. Optimizing Ultrametries with Infinity Norm
with S representing the set of n non-negative symmetric matrices with zero diagonal entries. Although O.zs is non- convex, this sub-problem admits a closed-form solution:
To compute the l ultrametric of a cost matrix D, we first construct its minimum spanning tree (MST) using Prim's ' algorithm (Prim, 1957), which yields an ultrametric tree T and the corresponding ultrametric Dr (Chen al., 2024). For simplicity, we collectively denote both T and Dr as MST(D), D). We then refine Dr to obtain the Lxx -neaeet ultrametric D using Theorem 5.
Case 1: If dij S. max{dak, dzk), the constraint is already satisfied, and D is feasible. Thus, DT = D E Pizs.
Case 2: If dg > (dxx. dy)), suppose du 2 djæ- To satisfy the constraint df S dfte- di must be decreased to drt, and du must be increased to d/.. Then the objective is
Theorem 5 -Nerrst Ultrametric (Chepoi & Fichet, 2000)). For a distance metric E I an optimal ultra- metric Dr. that minimizes Dr D1loo is given by:
Dr - D12 2 2. ((df) dijj² + -ddf -d+)))
D1. DS
+ 2 |MST(D) - D11
(12)
To minimize it, we set dt, dh. Let A = dij - dt, 2 0. then the objective becomes:
where 1; = 1 if * j., and 0 otherwise, for all j, E [m]
2 2A2 + (di - du -)))..
The algorithm is summarized below. Since Dř is an ultra- metric, the corresponding ultrametric tree T can be directly constructed using the MST. While prior works (Chen al., 2024; Lin et al., 2025) have used or compared against the MST (Prim, 1957), to the best of our knowledge, we are the first to leverage the loo -nearstt ultrametric for computing the tree- Wasserstein distance.
Minimizing this quadratic function gives A
1 (dl - dax).
df
A
dij
djk As shown above, we obtain the matrix Dr as follows:
de
Algorithm 1 UtrTaTWD-MST (Minimum Spanning Tree)
dy - - - A, di d/s = dik + A.
Input: DE I cost matrix.
Output: T": optimal ultrametric tree under infinity norm.
while all other elements of Dr remain unchanged from D. The resulting matrix Dr is denoted as Par,, (D), i.c. Parya (D) argmin |r DIl?, representing the pro- Dr jection of onto the subset Slikk
1: Compute the 1 -earest ultrametric via Eq. (12):
1 う |IMS((D) 5o 1.
Dj- = MST(D) +
2: Construct the ultrametric tree: T' MST(D?).
4

---

UWt: Optimizing Ultrametric Trees for Tree- Tree-Wascerstein Distance
Next, we solve min DrEn
DDr D. where Q is defined as:
3.5. Optimizing Ultrametries with Frobenius Norm:
Gradient Descent Method
:== (Dr E S" d. S max de.), Vi,jk E [n]]--
The UhraTWD-IP updates 2 entries in Dr E R= by projecting onto O(n²) subsets {Ryse), which can be inefficient for large n. To improve efficiency and scalability, we adopt a gradient descent (GD) approach to directly optimize the tree structure, inspired by Chen et al. (2024).
The feasible region represents the ultrametric space and can be expressed the intersection of all subsets: Ω Aaja The optimal solution is Pn(D), the projection of D onto Ωt but finding this projection directly is challenging.
To find an approximate solution, we employ an itera- tive projection method called the Halpern-Lions- Wittmann---n- Bauschke (HLWB) projection (Censor, 2006; Li et al., 2023). Starting with D D, we iteratively refine D by moving it toward D and successively projecting it onto each subset Payk. Empirically, each projection can enforce one strong triangle inequality, and we observe that iterative projections progressively satisfy more incqualities, poten- tially leading to an ultrametric. The iterative update rules are detailed below.
When the tree structure is fixed, the ultrametric Dr fully determined by the node heights in T, as defined in Eq. (8):
Dr = [dr (xt, xz)] [b(LCC((:,44)]]
which defines a mapping f Hr + Dr. where HT R denotes the heights of all 2n 1 nodes in T. By parameterizing Dr with Hr, the objective becomes:
F(Hr)= [I - D12
reducing the number of parameters from m(n-1) to 2n 1. 2
Theorem 6 (HLWB Projection (Censor, 2006; Censor et al., 2022)). Given subsets " 2m in Euclidean space with a non-empty intersection f2 = 52 n ... n.......... 9m and a distance matrix D E R initialize D = D. The iterative scheme for iteration I = 1, 2, ... is as follows:
The optimization process consists the following steps. Initialize T MDT(DD) with Hp. Then, for iteration t,
1. Update Hy: Fix the tree structure (i.e., the LCA rela- tionships) and update H via gradient descent:
H H1 oVF(H,
1. Movement: Compute Y o,D + (1 Dt)
2. Projection: Update D' = Ps, Pag Pam. (Y*),
where a is the learning rate. The gradient is efficiently computed using PyTorch's automatic differentiation.
where Pa,(Y (Y) denotes the projection Y onto subset Ωz 1000 is a steer-
If H1 are closed convex subsets and {at]t=)
2. Update Dr Compute D' = -HH)) Note that D' may have non-zero diagonal entries because d' corresponds to the height of leaf node 4, which can be non-zero in Hi. To ensure De is valid distance matrix, we adjust it as follows:
ing parameters sequence (i.e,, or E (0, 1) and o, + 0), D linearly converges to (D) as t- + cc.
The HLWB projection ensures convergence in convex prob- lems, but such guarantees may not hold in this non-convex case. Nevertheless, a high-quality approximate solution D' can be achieved in just one iteration (t = 1). This method updates matrix entries directly and implicitly alters the tree structure after each projection, as a different Dr induces different tree T. The explicit tree topology is constructed only once in the final step of Algorithm 2, using the MST algorithm on the projected matrix Dr. The algorithm is presented below, with EFI following Li et al. (2023).
1 d' + - (2d2 di d'g))
which ensures zero diagonal values.
3. Update T: Apply the MST algorithm to D' to update the tree structure 7 and obtain new heights H7. Then repeat the update of H4, in Step 1.
While based on Ultra (Chen et al., 2024), our method uses an unsupervised objective |r D1I: instead of a supervised loss, resulting in a more efficient Algorithm 3.
Algorithm 2 UltraTWD-IP (Iterative Projection)
Algorithm 3 UltraTWD-GD (Gradient Descent)
Input: E R"**; cost matrix, m: maximum number of iterations (default m = 1).
Input: D E R=x,, cost matrix, m: maximum iterations
(default m = 8), a: learning rate (default a = 0.02).
Output: T": the ultrametric tree under the Frobenius norm.
Output: T the ultrametric tree under the Frobenius norm
1: Initialize D D.
1: Initialize 7 MST(D) with node heights H..
2: for (or1=1 1 to do
2: for t to m do
1 D + D 3: Update D' +- r+1
(Movement step)
3: Update heights: 11 + -H VVF(
4: for each triplet ki,j,k) do
4: Compute the ultrametric: D' = f(Ht)
5: D' + P (D')
(Projection step)
5: Adjust entries of DC: d' < } (2d' dr - d'
6: end for
6: Update tree: 7 = MST(D') with heights lf.
7: end for
end for
8: Construct the ultrametric tree: T MST(D*).
Return the optimized ultrametric tree: T = T
5

---

UWD Optimizing Ultrametric Trees for Pree-Wasscrstein Distance
3.6. Algorithm Analysis
Table 2. Statistical Information of Text Datasets.
Convergence Analysis. Due to the non-convex and NP-hard nature of the problem (Eq. (11)), theoretical convergence is extremely hard to guarantee. Similar gradient descent methods in prior work (Chierchia & Perret, 2019; Chen et al., 2024) also lack convergence guarantees. Neverthe- less, both UltraTWD-IP and GD methods exhibit empirical convergence, as demonstrated in Fig. 4.
Dataset
Avg. Words (x)
# Test Data (.V
# Classes
BBCSport
6,051
220
Reuters
6,416
1,000
8
Ohsumed
9,467
1,000
10
Recipe
4,084
1,311
15
Implementation. Each dataset contains 5 test sets. For each test set, we construct the vocabulary X [i....., un. Rdx=, where d' = 300 is the word embedding dimension and v is the number of unique words (average values shown in Table 2). Each test document f4; is represented as a normalized bag-of-words distribution. We compute the cost matrix D E R=*= using Euclidean distances: de = ||æ; æy The vocabulary X and cost D are then used to learn the tree T,, and each WT(Ps, M4) is computed based on T.
Time Complexity of Tree Learning. Given a fixed support size n. both Ult and UltraTree learn single tree. (1) UltraTWD- Total complexity is (n²)), as computing the MST for D E R takes (n²²) time. (2) UltraTWD- IP: It involves ((n²) projections with (1) cost each. We use only one iteration in practice, resulting in ((n²) total complexity, which is practical for moderately large n. (3) UttraTWD-GD: Each iteration takes O(n") due to gradient descent and MST, giving (mm²) overall. We set 8 in practice, making it more scalable than UltraTWD-IP for large n. (4) UltraTree (Chen et al., 2024): Precomputing W1 for M training pairs requires O log n). Tree learn- ing has O(mn² + m) complexity due to computing Wr in each iteration for loss evaluation. Overall, UlrrTWDD-G significantly more efficient than UltraTree.
Evaluation. The effectiveness of the tree- Wasserstein dis- tance is assessed based on the following objectives: (I) Approximation Accuracy: Measuring the error between tree- Wasserstein distance and W-Wasserstein distance.
(2) Application Performance: Evaluating the utility of tree- Wasserstein distance in text-based tasks, including docu- ment retrieval, ranking, and classification.
(3) Performance Analysis: Analyzing the impact of hyper- parameters and the algorithm efficiency trade-off.
Table 1. Time Complexity Comparison. For distributions over
support of size n. "W Computation" is the cost of computing the pairwise distance matrix Wr or W1 E T N
Baseline Methods. We compare the UWD methods
Time Complexity
Tree Learning
W Computation
with 10 approximation methods across various categories. Detailed descriptions are provided in Appendix A.4.
O(1²)
O(N n)
UI-II
O(e))
O(N2 n)
(1) Entropy-based method: Sinkhorn distance (Cuturi,
UDrGTWDDGG
O(mm²)
O(N2. n)
2013) is implemented using the POT library (Flamary et al., 2021), with a regularization parameter X = 1 and a maxi- mum of 100 iterations.
UltraTree
O(mn² 2 + mm
O(N2. n)
Wi
O( n²² logn)
(2) Tree-construction methods: QuadTree (Indyk & Thaper,
Time Complexity of W Computation. For N distribu- tions over support of size n, computing the 1-Wasserstein distance matrix W1 E N requires O(N2. n log n) time, while computing the WT matrix on ultrametric tree requires only O .n) time. For a fair comparison, includ- ing UltraTWD's D's tree learning cost (On')) ) or (mn))), our methods remain significantly faster than exact W1 compu- tation, especially for large datasets. Results in Section 4.4 further demonstrate the computational efficiency of TWD.
2003) and ClusterTree (Le et al., 2019), implemented by
modifying the code from Yamada et al. (2022).
(3) Weight-optimized methods: Weight-optimized QuadTree
(TDD)) and ClusterTree (CTWD) with regularization parameter A = 0.001 (Yamada et al., 2022).²
(4) Tree-sliced methods: Sliced -eaeed--uaTTree Sliced- ClusterTree, Sliced-qTWD, and -li, average
TWD over 3 randomly sampled trees (Yamada et al., 2022). (5) Supervised method: UltraTree (Chen et al., 2024) is trained on 1,000 randomly generated distributions with 1 1% sparsity and precomputed 1-Wasserstein distances.'
4. Experiments
4.1. Experimental Setting
All experiments were conducted using Python 3.8 on a Linux server with an AMD EPYC 7742 64-Core Processor, 256 logical CPUs, and 512 GB RAM. The average perfor- mance and standard deviation are reported across the test sets. Additional details can be found in Appendix B.
Dataset. We evaluate our methods using four benchmark text datasets: BBCSport, Reuters, Ohsumed, and Recipe (Huang et al., 2016), following previous studies (Takezawa et al., 2021; Chen et al., 2024). These datasets are publicly available with detailed statisties provided in Table 2.
FFolowing Chen al. (2024) to balance runtime and accuracy. hittps://githube com/oi st/treeOT com/ httm//hhhe///eeeeeeeeeeeeee 1earnenn__
hhttps://git thub. com /wwuner/wdddd
6

---

UWt: Optimizing Ultrametric Trees for Tree- Tree-Wasserstein Distance
Table Comprehensive comparison of tree-Wasseee Wasserstein distance methods across four text datasets. Bold indicates the hest result,
underline marks the second -bstt, and wavy underline denotes the third-best. Our UttraWD-IIP and GD methods outperform the baseline methods in most cases across (a) approximation error, (b) document retrieval, (c) document ranking, and (d) document classification.
Metric
(a) RE-W J: Approximation Error of Wr
(b) Precision 1: Document Retrieval
Dataset
BBCSport
Reuters
Ohsumed
Recipe
BBCSport
Reuters
Ohsumed
Recipe
Sinkhom
0.258 1:0000
0.167
0.166 0.002
0.244 4-0.00
0.4412+00u 4..09
0.207= 4.354
QuadTrec
0.643 +-0080
0.630
0.649 -0,089
0.745 +-0.09
0.722-+0.026
0.732++0015
0,552
0.694
ClusterTree
0.733 0.774
0.768_0.0m
0.736110.00
0.548
0.543
0.315
0.395
qTWD
0.162
0.
0.139
0.8 0.810
0.81940000
0.69 1 10000
0.7 0.795
cTWD
0.11610000
0.079
0.097
000004
0. 0..000
0.8522000004
0.840 002
0.723
0.820 20.002
Bieedd-uudTTree
0.689 4:0042
0.651 #0001
0.681 -0.034
0.758 +0.01
0.795
0.801 +0.004
0.655
0.769 20.005
0.734
0.770 0.004
0.747 +0.00
0,637 +0.015
0.652 oox
D, 0.445
0.522
-lD
0.104
0. 40.001
0.127=0:00
0. t..00
0,82 +0.00s
0.824
0.708
0,807
Sliced-eIWD
0.116::0002
0.079910001
0.095_0.000
0. 10.000
0,863 1.0002
0.849 1:0000
0.742.10.005
0.83 0.000
UltraTree
+0000
0,041 +00003
0.038-0.001
0.0227+0.00
0.000
0. 0.834
0.749 +0.003
0.830 200004
0.109 100001
0.050.00000
0.077
-0:004
0.098 +00007
0. 0.860
0,854
0,765++..00
0,840
00.00s
UlraTWD-IIP
0.014 100000
0.033310.00
0.036 0.001
0.026 1.0000
0.885 0,0000
0.876 100000
0.788 0.000
0.866 00,00
UtraTWD-GD
0.028 00000
0.0.30+0.00
0.017 =0,000
0.035
0.867
0.860+00002
0.776
0.848 0.0002
Metric
(c) MRR t: Document Ranking Ohsumed
(d) Accuracy t: Document Classification
Dataset
BBCSport
Reuters
Recipe
BBCSport
Reuters
Ohsumed
Recipe
Sinkhorn
0.533 +-0.017
00.895
0.395
0.226-10.01
0.359 +00000
0.773
0.2924
0.381=
QuadTree
0.78
0. 0.790
0.662_000%
0.761 10:06
0. 0.805=0.010
0.88940014
0.398
0.488
ClusterTree
0.588
0.603
0.377
0,452++0.2
0.77 20.041
0.891
0. 0.40642077
0,461 20.007
qTWD
0.865.0000
0.865
0.783_
0.8522.001
0.819=0.020
0.892+0011
0.394
0.493
cTWD
0.884
+..07
-..10
0.873 +0.0
+..022
0.896
0,41 I +0026
0.493=00.00
Sliced- Quad'Tree
0.846
0.848
40.099
0.754_0.013
0.827 +0.00n 0..07
0. 4..23
0.892+00055
D. 40.22
00.07
Sliced- ClusterTree
0,676
0,715
0,515
0,585
0. 0.798
0,902++0009
0,423 +0007
0,476
Sliced-qTWD
0.872:20001
0.869 10.006
0.796_
0.862
0.8 0.819
0.893 10013
0.396_10028
0.494
Sliced-CTWD
0.890
10.005
0.819 00,10
0.880.10.00
0.81 0.817
0.897 10013
2.0026
0.494=00.00
UltraTree
0.890
00.00
0,83 00000
0.877+00012
0.826+0.017
0.89440014
0.495 =0.007
UltaTWD-MMS
0,909 +.0008
0.893.40.005
0,845
-0:008
0,886
0.835
0.899400r
0.421 1 +0024
0,495 000007
UhraTWDDIIP
0.924 100011
0.913.10.005
0.863 0.005
0.908
10.001
0.839 10.017
0.905 10.000
0.428 10.020
0.495 0.000
UtraTWD-GD
0.922 1:0.009
0.899:10.006
0.855
0.010 0,010
0.895 100007
0.836
0.899 10003
0.427
0.495 0.007
4.2. Approximation Error of Free-Wasserstein Distance
(2) Tree-construction Method: Quad'Tree and ClusterTree exhibit large approximation errors (e.g., 0.643 and 0.774 on the BBCSport dataset) because their edge weights are fixed based solely on tree depth, resulting in suboptimal tree representations and poor accuracy.
Each test set contains N documents {,I,..., | repre- sented as normalized ofg-of--wrrd distributions. We com- pute the pairwise 1- I-Wasserstein distance matrix W1 = [Wi [k)] E N as the ground truth and the tree- Wasserstein distance matrix W7 = [Wr(M, My)] E R N as its approximation (including the Sinkhorn method). We evaluate the Relative Error of Wr (RE-W), defined as:
The qTWD and cTWD methods enhance QuadTree and ClusterTree by optimizing edge weights, significantly reducing the relative error on the BBCSport dataset from 0.643 (QuadTree) and 0.774 (Clus- terTree) to 0.162 (QTDD) and 0. 16 (cTWD), respectively. (4) Tree-sliced Method: Sliced methods average TWDs over multiple random trees, yielding comparable or slightly higher errors due to the variance from multiple trees.
(3) Weight-optimized Method:
|Wr Wille
RE-W-
|Will
Table 3(a) shows that our UltraTWD-IP approach achieves the smallest RE-W compared to all haseline methods.
Overall, unsupervised baselines cannot optimize tree struc- tures and instead rely on randomly constructed trees from Quad or ClusterTree, resulting in suboptimal accuracy. In contrast, our methods optimize both the tree structure and edge weights by refining the ultrametric Dr. Up to iso- morphism, each Dr corresponds to unique tree structure with specific edge weights. By aligning Dr closely with the cost matrix D, all three UlrraWW methods dynamically adjust the tree structure, achieving lower errors across all four datasets compared to unsupervised bascline methods.
Comparison with Unsupervised Methods:
(1) Entropy-based Method: Sinkhorn distance exhibits higher errors due to entropy regularization, which improves speed at the cost of accuracy. A smaller regularization pa- rameter À and more iterations can improve accuracy but significantly slow down the computation. For example, on the Recipe dataset, setting À = 0.01 with 300 iterations re- duces RE-W W from 0.244 to 0.002, but runtime rises from 0.1 to 2.2 hours, while UTrrIP completes in 1.2 hours.
7

---

UWD: Optimizing Ultrametric Trees for Tree-Wasserstein Distance
Table 4. Approximation error Utra and Utt methods.
Document Classification. For text classification, we compute the kernel K exp(-Wr//o where median(Wr) and apply a kernel SVM with 10-fold cross- validation. As shown in Table 3(d), UltraTWD-B achieves the highest classification accuracy across all four datasets, showing its practical advantage for downstream tasks.
RF-W- = Wy--Wale and RE-D Dlle,,
Bold is the best
[Will
D
result and underline is the second-best.
Dataset
Reuters
Ohsumed
Metric
RE-W
RE-D
RE-W
RE-D
Tree-T%
0.041:0000
0.113
0.038
0.12340000
-ree27e
0.061
0.130
0.069
0. 0.147
UltraTree-100
0.187 0.127
0.119
0.18940000
Ulra
0.033
0.102
0.036
0.1 15 +0001
4.4. Performance Analysis
UltaaWD--G
0.030 0.100
0.017 10.000 0.101 +0,000
Hyperparameter and Convergence Analysis. Fig. 4 shows the impact of iterations and learning rates on RE-D convergence. Both UlraTWD-IP and GD methods demon- strate empirical convergence, though not necessarily to a global minimum. (1) UltraTWD-IP: It converges in fewer iterations and achieves higher retrieval precision. In prac- tice, one iteration is often sufficient. (2) UltraTWD-GD: A learning rate of 0.05 accelerates RE-D reduction but may overshoot local minima, leading to unstable precision. A learning rate of 0.02 provides stable convergence and high precision, and used in all experiments. Additional results are provided in Appendix C.2.
Comparison with Supervised Methods: UltraTree, a state-of-the-art baseline, learns the ultrametric Dr by mini- mizing the loss || ||Waaann W [ on training data. How- ever, its performance is highly sensitive to the sparsity of training distributions, which limits its robustness. As shown in Table 4, the approximation error is low at 1% sparsity, as it matches the average sparsity in the test data. When sparsity increases to 2% or 10%, the RE-W on the Reuters dataset rises from 0.041 to 0.061 and 0.127, respectively, along with higher RE-D. This indicates that UltraTTree strug- gles to generalize when the sparsity levels between training and test distributions differ.
MST - we GD (a=0.01)
GD (a=0.02)
GD (a=0.05)
In comparison, our UltraTWD-IP and GD methods are in- herently robust to sparsity, as they do not require training data. Instead of minimizing || ||Wyain W W p we mini- mize the fundamental difference |D DIf resulting in lower RE-D and RE-W compared to UltraTreee
0.3
0.89
RE-D
0.2
Precision
0.88
0.87
0.1
0.86
0.85
0.0
0
5
10
15
20
0
5
10
15
20
4.3. Applications of
Distance
Iteration
Iteration t
Figure 4. Hyperparameter analysis of UtrraWWD methods on the Reuters dataset, illustrating empirical convergence.
We further evaluate the quality of tree- Wasserstein distance matrices in document retrieval, ranking, and classification, showing the utility of our methods in text applications.
Efficiency Analysis of Learning Time. Table reports the average time for tree learning. UtraTWD-MST's is the fastest (as fast as ClusterTree), completing in just 2-7 seconds, while achieving much lower errors than all unsupervised baselines. UltraTWD- GD offers a great balance between ef- ficiency and accuracy, with learning times of 15-72 seconds, comparable to Biced-gTWDDeTWD. but achieving much smaller errors and higher retrieval precision. Compared to UltraTreee, UltraTWD-GGD is up to 43x faster and also more accurate, making it well-suited for real-world applications.
Document Retrieval. We evaluate performance on a nearest neighbor retrieval task. For each test document, the top- 10 nearest neighbors among the remaining N - 1 documents are retrieved using both WT and Wi, resulting in two sets: IT and 41, respectively. We use Precision as the evaluation metric, defined as Precision |rr//| that our methods achieve the highest precision, following the order: UltraTWD-IP > UWtr-TDDGG > UltraTW-- MST > most baselines. This demonstrates that UltraTWWD effectively preserves neighborhood structures, ensuring that similar documents under W, remain close under Wr. Visu- alizations are provided in Appendix C.1.
10 Table 3(b) shows
Efficiency Analysis of W Computation. We generate 100 random distributions over the BBCSport vocabulary (average 6,051 words), each containing R-aaldd words (non- zero entries). We compare the total time required to compute the matrix W₁ or WT E R 1000x100 including tree learning time, As shown in Table 6, when rivlid 5, 000, computing the W1 matrix takes 5.0 hours, with each Wi (p) requiring 3.6 seconds. In contrast, UttrTWW variants complete the Wr matrix within hour, with each (T(, v) computed in just 0.5 seconds, achieving a significant speedup. UlraTres is slower than UltraWWD due to its costly training process, including the precomputation of W, for training pairs.
Document Ranking. We assess ranking performance using the Mean Reciprocal Rank (MRR), calculated as MRR
1 N Em rank, For a given query document 14, rank; is the position of the first relevant item in WT based on the ranking induced by W1. MRR reflects the average ranking quality across N queries, ranging from (worst) to 1 (best). Table 3(c) shows that our methods achieve the highest MRR, with relevant items consistently ranked near the top.
1
8

---

UDt Optimizing Ultrametric Trees for Tree- Tree-Wasserstein Distance
Table Tree learning time (in seconds) TWD methods.
Bake, Fry
Apple, Watermelon
Chicken, Beef
Lettuce
Dataset
BBCSport
Reuters
Ohsumed
Recipe
QuadTree & qTWD
ClusterTree & cTWD
UltraTWD
Average Words (n)
6,051
6,416
9,467
4,084
QuadTree
840
940
1340
6
ClusterTree
3
340
540
2+0
gTWD
14
14+2
20+2
1042
Bake
Fry
Lettuce
Bake
Fry
cTWD
13
13:12
15:1
11:1
Chicken
Boef
Lettuce
rieed-QuadTree
27.11
27:11
41
16:10
Bake
Chicken
Apple
Watermelon Fry
Beef
Lettuce
Apple
Waternelon
Chicken
Beef
Apple
Watermelong
Sliced-ClusterTre
10
10.00
1641
6
Sliced-qTWI)
4022
3941
5543
28
Sliced-cIWD
37
3743
4242
3141
UltraTree
4217
938,29
1494
361
(b)
U-raTWD-MST
340
340
740
2+0
UWDITTWD-IP
1207=77
1437
45834s8
37344
Figure 5. Tree structures produced by UltraWWD baseline
UhraTWD-GD
28
34
72+21
1543
methods for 7 semantically grouped words from the Recipe dataset.
Table Total time (in hours) of computing Wi or Wr E
Acknowledgments
R 1000 100 the BBCSport dataset, including tree learning time.
Moaid represents the number valid words in each distribution.
The work of Fangchen Yu, Jianfeng Mao and Wenye Li was supported in part by Shenzhen Science and Technology Pro- 02222002209900000000000000
Valid Words (Nvmss
1,000
2,000
3,000
4,000
5,000
Wi
0.2
0.9
2.3
3.6
5.0
in part by the Na- tional Natural Science Foundation of China (No. 72394362 and U1733102) 102). in part by the Guangdong Provincial Key Laboratory of Big Data Computing, The Chinese University of Hong Kong, Shenzhen (No. 2024SC0003), in part by the Shenzhen Science and Technology Innovation Committee under the Shenzhen Stability Science Program grant (No. 2024SC0010), in part by the Shenzhen Research Institute of Big Data (No. J001 20210005001) and in part by CUHK- Shenzhen (No. PF.01 .000404). The work of Fangchen Yu and Qiang Sun was supported in part by the Natural Sci- ences and Engincering Research Council of Canada (Grant SGPIN-2018-0648A computing resources provided by the Digital Research Alliance of Canada, and MBZUAI.
gram (No.
UltraTree
0.8
1.0
UlraTWD-MST
0.2
0.4
0.5
0.6
0.7
UlrrTTWD-IP
0.6
0.7
0.9
1.0
1.0
UlraTWD-GD
0.2
0.4
0.5
0.6
0.7
4.5. Qualitative Comparison of Tree Structures
To qualitatively assess the semantic interpretability of the learned trees, we select 7 representative words from the Recipe dataset, grouped into 4 semantic categories: Cook- ing (Bake, Fry), Fruit (Apple, Watermelon), Meat (Chicken, Beef). and Vegetable (Lettuce). The word embedding matrix X 7x300 and the Euclidean distance matrix D E 7 structures. used as input. Fig. presents the resulting (a) QuadTree & gTWD produce flat and shallow trees with- out a semantic structure. (b) ClusterTree & cTWD partially group related words but often mix categories. In contrast, (c) UTD clearly separates all four semantic groups and produces a well-structured hierarchy. In this simple case, all three UltraTWD variants yield the same tree, demonstrating UltrTWD's strength capturing semantically meaningful structures that reflect underlying word relationships.
Impact Statement
This paper presents UltraTWD, novel framework for tree- Wasserstein distance that aims to advance the field of ma- chine learning by improving the efficiency and accuracy of 1- I-Wasserstein distance approximations. The proposed methods have the potential to benefit wide range of appli- cations, including text retrieval, ranking, and classification tasks, by offering accurate and reliable solutions.
5. Conclusion
Overall, we believe UltraTWD contributes positively to the field of machine learning and offers practical benefits across various domains, with no immediate ethical concerns requiring specific attention.
We introduce UltraT'WD a novel unsupervised framework for tree-Was Sre--Waserssten distances that jointly optimizes the tree structure and edge weights, addressing the limitations of existing methods. By solving ultrametric nearness problems, develop algorithms leveraging minimum spanning trees, iterative projection, and gradient descent. These meth- ods demonstrate superior accuracy in I-Wasserstein distance estimation and significantly improve performance tasks such as document retrieval, ranking, and classification, mak- ing them valuable for Wasserstein-based applications.

---

UltraTWD: Optimizing Ultrametric Trees for Tree- Tree-Wascerstein Distance
References
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems, 26, 2013.
Agarwal, P., Raghvendra, S., Shirzadian, P., and Sowle, R. A higher precision algorithm for computing the 1- wasserstein distance. In International Conference on Learning Representation, 2023.
Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z.,, Bois-
bunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras, K., Fournier, N., Gautheron, L.., Gayraud, N. T., Janati, H., Rakotomamonjy. A., Redko, I., Rolet, A., Schutz, A., Seguy. V., Sutherland, D. J., Tavenard, R., Tong, A., and Vayer, T. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021.
Agarwal, P. K., Raghvendra, S., Shirzadian, P., and Yao, K. Fast and accurate approximations of the optimal transport in
and discrete settings. In Proceedings of the 2024 Annual ACM-SIAM Symposium on Discrete Algorithms, pp. 4514 4529. SIAM, 2024.
Ailon, N. and Charikar, M. Fitting tree metrics: Hierarchical clustering and phylogeny. SIAM Journal on Computing, 1092) 2011.
Fox, E. and Lu, A deterministic near-linear time approx- imation scheme for geometric transportation. In IEEE 64th Annual Symposium on Foundations of Computer Science, pp. 1301-1315, IEEE, 2023.
Backurs, A., Dong, Y., Indyk, P., Razenshteyn, I., and Wag- ner, T. Scalable nearest neighbor search optimal trans- port. In International Conference on Machine Learning, pp. 497-506. PMLR, 2020.
Gavryushkin, A. and Drummond, A. The space of ultra- metric phylogenetic trees. Journal of Theoretical Biology, 403:197-200 2016.
Censor, Y. Computational acceleration of projection algo- rithms for the linear best approximation problem. Linear Algebra and Its Applications, 416(1):111-12)) 2006.
Gonzalez, T. F. Clustering minimize the maximum in- tercluster distance. Theoretical Computer Science, 38: 293-306, 1985.
Censor, Y., Moursi, W. M., Weames, T., and Wolkow- icz, H. Regularized nonsmooth newton algorithms for best approximation with applications. arXiv preprint 1rXiv:2212.13182, 2022.
Huang, G., Guo, C., Kusner, M. J., Sun, Y., Sha, F., and Weinberger, K. Q. Supervised word mover's s distance. Advances in Neural Information Processing Systems, 29, 2016.
Chatziafratis, V., Niazadeh, R., and Charikar, M. Hierarchi- cal clustering with structural constraints. In International Conference Machine Learning, pp. 774-783. PMLR, 2018.
Indyk, P. A near linear time constant factor approximation for euclidean bichromatic matching (cost). In Proceed- ings of the Eighteenth Annual ACM-SIAM Symposium on Discrete Algorithms, volume 7, pp. 39-42, 2007.
Chen, S., Tabaghi, P., and Wang, Y. Learning ultrametric trees for optimal transport regression. In Proceedings of the AAAI Conference on Artificial Intelligence, vol- ume 38, pp. 1005572006655 2024.
Indyk, and Thaper, N. Fast image retrieval via embed- dings. Workshop on Statistical and Computational Theories of Vision ICCV, 2003.
Chen, X., Jayaram, R., Levi, A., and Waingarten, E. New streaming algorithms for high dimensional emd and mst. In Proceedings of the 54th Annual ACM SIGACT Sympo- sium on Theory of Computing, pp. 222-233, 2022.
Khesin, A. B., Nikolov, A., and Paramonov, D. Precondi- tioning for the geometric transportation problem. Journal of Computational Geometry, 11(2):234 259, 2021.
Khrulkov, V., Ryzhakov, G., Chertkov, A., and Oseledets, I. Understanding ddpm latent codes through optimal transport. In The Eleventh International Conference on Learning Representations, 2023.
Chepoi, V. and Fichet, B. loc lppappproinnttion via subdomi- nants. Journal of Mathematical Psychology, 44(4):600~ 616, 2000.
Chierchia, G. and Perret, B. Ultrametric fitting by gradient descent. Advances in Neural Information Processing Systems, 32, 2019.
Křivánek, M. The complexity of ultrametric partitions on graphs. Information Processing Letters, 27(5):265-270, 1988.
Cohen-Addad, V., Kanada, V., Mallmann-Trenn, F., and Mathicu, C. Hierarchical clustering: objective functions and algorithms. In Proceedings of the Twenty-Ninth An- mual ACM-SIAM Symposium on Discrete Algorithms, pp. 378-397, 2018.
Kusner, M., Sun, Y., Kolkin, N., and Weinberger, K. From word embeddings to document distances. In International Conference on Machine Learning, pp. 957-966. PMLR, 2015.
10

---

UWt: Optimizing Ultrametric Trees for Tree- Tree-Wasserstein Distance
Le, T., Yamada, M., Fukumizu, K., and Cuturi, M. Tree- sliced variants of wasserstein distances. Advances in Neural Information Processing Systems, 32, 2019.
Yu, E,, Bao, R., Mao, J., and Li, W. Highly-efficient robinson-foulds distance estimation with matrix correc- tion. In 26th European Conference on Artificial Intelli- gence, pp. 2914-2921 IOS Press, 2023.
Li, W. Yu, F., and Ma. Z. Metric nearness made practi- cal. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 8648-8656, 2023.
Lin, Y.-W. E., Coifman, R., Mishne, G., and Talmon, R. wree-wasserstein distance for high dimensional data with latent feature hierarchy. In The Thirteenth International Conference Learning Representations, 2025.
Otao, S. and Yamada, M. A linear time approximation of wasserstein distance with word embedding selection. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 15121- 15134, 2023.
Pele, O. and Werman, M. Fast and robust earth mover's distances. In 2009 IEEE 12th International Conference on Computer Vision, pp. 460-467. IEEE, 2009.
Prim, R. C. Shortest connection networks and gen- cralizations. The Bell System Technical Journal, 36(6): 1389-1401, 1957.
Rabin, J., Peyré, G., Delon, J., and Bernot, M. Wasserstein barycenter and its application to texture mixing. In Scale Space and Variational Methods in Computer Vision, pp. 435-446 Springer, 2012.
Semple, C. and Steel, M. Phylogenetics, volume 24. Oxford University Press Demand, 2003.
Takezawa, Y., Sato, R., and Yamada, M. Supervised tree- wasserstein distance. In International Conference on Machine Learning, pp. 10086-10095. PMLR, 2021.
Tran, H. V., Nguyen, K. N., Pham, T., Chu, T. Le, T., and Nguyen, T. M. Distance based tree-sliced wasserstein distance. In The Thirteenth International Conference on Learning Representations, 2025.
Tran, V.-H., Pham, T., Tran, T., Le, T., and Nguyen, M. sliced wasserstein distance on system of lines. arXiv preprini arXiv:2406.1 1at.222466.375.... 2024.
Villani, C. Optimal Transport: Old and New, volume 338.
Springer, 2008
Wang, D. and Wang, Y. An improved cost function for hierarchical cluster trees. Journal of Computational Ge- ometry, 11())283-331, 2020.
Yamada, M., Takezawa, Y., Sato, R., Bao, H., Kozareva, Z., and Ravi, S. Approximating 1-wasserstein distance with trees. Transactions on Machine Learning Research, 2022.

---

Ult: Optimizing Ultrametric Trees for Tree- Distance
A. Background and Related Work
A.2. Background of Tree-Wassersteil Distance
A.I. Background of Wasserstein Distance
Distance. The 1- Wasserstein distance requires solving linear programming problem to obtain the optimal transport plan y which can be computationally expensive. To improve the efficiency, the tree- Wasserssnnt distance encodes the cost d( 2z)) into the shortest path distance (z,, 25)) defined on the tree T. Then the tree- Wasserstein distance defined on the tree has a closed-form solution:
Motivation of
The Wasserstein distance (Villani, 2008) measures the distance between probability distributions. Consider two discrete measures Jr = Σi=1 agd2, and = Eÿ= by represented 1-dimensional histograms shown below:
wp(p,))
bj
Wr (14, v) = Eww (m((T)) - vTTe).
= TET
(13)
where we is the weight of each edge in the tree. The derivation of this formula can be found in Proposition I of Le et al. (2019).
V
The p- o-Wasserstein distance is defined as
1/p
Linear Time Complexity. The closed-form solution given in Eq. (13) requires summing over all edges in the tree. In a rooted tree, the number of edges is equal to the number of nodes, denoted by Temade- Therefore, the computation of the trss-eeann distance has a time complexity of O (n Since the number of nodes is bounded by (()), where represents the number of leaf nodes, the overall complexity is also (n) (Chen et al., 2024). Consequently, the tree- Wasserstein distance can be computed in linear time with respect to the number of leaf nodes.
v)(p,v) =
min (pE()
Σ (==1
Yey Te)d( xz))
where y is the transport plan in the set [[ v) defined
((,,2) (7 E R R"*" Σn az Σ Ti = bj. Vi.j}.
In this paper, we focus on the commonly used 1- I-Wasserstein distance, defined as
Wi (1 = min myd((;;,))) yEl(p.)) KJ=1
A.3. Background of Ultrametric
which is also known as the optimal transport distance or Earth Mover's Distance (Pele& & Werman, 2009).
Definition of Ultrametric. An ultrametric D E Rax= is a special type of tree metric that satisfies the strong trian- gle inequality: di S max(dik, d for all i, k. This implies that any triplet (i.j.k) forms an isosceles triangle. Ultrametries naturally correspond to rooted trees with n leaf nodes, where all leaves are equidistant from the root. In such trees, the distance between two leaves is defined by the height of their lowest common ancestor. The resulting pairwise distance matrix among the leaves is an ultrametric matrix D. Note that ultrametric trees are rooted but not necessarily binary. In practice, we follow Chen al. (2024) and use the hierarchical minimum spanning tree procedure (Prim, 1957) to construct binary ultrametric trees.
At a high level, the cost d( Rÿ) represents the cost of transporting unit mass from H; to Zj. Thus, the 1- Wasserstein distance computes the minimum cost required to transport all mass from {a,} to {by}, where "y specifies the amount of mass moved from 2; to Z.j.
In text analysis, x; denotes the embedding vector of the i-th word in the vocabulary, and the vocabulary of n words represented by a word embedding matrix = [x1, T2...., 2n] E Rdx,, document μ is modeled as normalized bag-of-words distribution over the vocabu- lary X, expressed as μ Σ/=1 a;da,, where a, is the fre- quency of the i-th word in the document. The normalization (E ag = 1) ensures μ is a valid probability distribution.
Benefits of Ultrametrics. Ultrametrics are preferred over general tree metrics in our setting for several reasons. First, constructing a tree from a general tree metric is nontrivial, while ultrametric trees can be efficiently generated using algorithms such as the minimum spanning tree. Second, ultrametries satisfy the strong triangle inequality, which is easier to enforce and analyze compared to the more com- plex four-poin condition required for general tree metrics. Finally, ultrametrics are widely used in practice, including hierarchical clustering (Ailon & Charikar, 2011) 1) and phylo- genetic tree construction in bioinformatics (Gavryushkin Drummond, 2016; Yu et al., 2023).
When the Euclidean distance between word vectors is used as the cost, i.e., d(x,, .,)) = æa - 2t ||2, the I-Wasserstein distance corresponds to the Word Mover's Distance (WMD) (Kusner et al., 2015). Intuitively, WMD measures the mini- mum cost needed to transport the words from one document to another.
Despite its effectiveness, the Wasssersseein distance has a computational complexity of O(n² log n), making it com- putationally expensive in practice (Villani, 2008).
12

---

UWD: Optimizing Ultrametric Trees for Tree-Wassesstein Distance
Table Comparison of UTD with existing Wasserstein distance methods. The table evaluates four key aspects: (1) the ability to optimize the tree structure, (2) the ability to optimize edge weights, (3) the dependency on multiple trees, (4) the need for embedded data points, and (5) need for training data. Unlike traditional methods, UltraIWD achieves both tree and edge weight optimization while avoiding the dependency multiple trees, embedded data points and training data, making it more accurate and robust.
Method
Tree Structure
Edge Weight Optimization
No Multiple
No Embedded
No Training
Optimization
Trees Needed
Data Points Needed
Data Needed
QuadTree
X
X
ClusterTree
X
qTWD
X
cTWD
X
Sliced-QuadTree
X
X
Sliced-ClusterTree
X
X
X
Sliced-gTWD
X
X
Sliced-eTWD
X
X
UltraTree
UTrTTWD-MMT
UlraTWD--IP
UltaTWD-GD
A.4. Comparison Methods and Limitations
partitioning methods, this algorithm iteratively divides the data into predefined number of clusters by apply- ing the farthest-point clustering strategy. The centroid of each cluster is designated as node the tree, and the partitioning process is recursively repeated for each cluster until the desired depth is reached.
Our approach is evaluated against a range of representative methods designed to approximate I-Wasserstein distance:
Entropy-Based Method:
The Sinkhorn distance (Cu-
turi, 2013) introduces an entropic regularization term to the objective function to improve the efficiency:
Once the QuadTreee or ClusterTree constructed, the weight w. of each edge the is determined by its depth ((e), which represents the number of edges in the unique path from the root to the edge e. The edge depth reflects the hierarchical level of the edge, with larger values indicating deeper levels in the hierarchy. Based on this depth, the edge weight is defined as we = 22 g(((),, leading to the tree- Wasserstein distance formula:
Wx(,
min (EL()M)
Σ hydla; XH()), (14) 1J=1
where À > 0 is the regularization parameter and (+)) is the entropy of the transport plan defined as:
H(7) = Σ
Thi log 7ij-
The Sinkhorn distance can be efficiently computed using the Sinkhorn- Knopp algorithm with the computational com- plexity of (n²) (Cuturi, 2013).
WT(,v) 2 (I(TT)) -v(Te) eET
II. Tree-Coree Construction Methods: They build hierarchical tree using either QuadTree or ClusterTree structures.
III. Weight-Optimized Methods: The TWD and cTWD methods (Yamada et al., 2022) optimize the edge weights of fixed QuadTree or ClusterTree structure by solving a Lasso-based regression problem:
QuadTree (Indyk & Thaper, 2003; Backurs et al., 2020) constructs a hierarchical tree by recursively par- titioning the data space, represented as a hypercube in Rd, into smaller uub-hypereubes. The process begins with an initial hypercube, which iteratively divided into 2d sub-hypercubes. If sub-hypereube contains exactly one data point, its center is used as node in the tree. If it contains multiple data points, the partitioning continues until either predefined depth is reached or no further subdivision is required. Finally, this method constructs tree by splitting a space into hypercubes.
min w E. "rande
Σ (d(#;, wT ==.5)²+ A||w||, (15) 1-==1
where w E R.iiee represents the vector of edge weights, and Zi,j E Riits is vector determined by the tree struc- ture. In this formulation, the tree distance is expressed as dm m)) w² zi.j. The optimization problem aims to adjust the tree distance dr to align with the cost d. The l -rregllarizztton term encourages sparsity by driving some edge weights to zero, allowing the corresponding nodes to be merged, thereby reducing the overall tree size.
ClusterTree (Le et al., 2019) adaptively partitions the data space using farthest-pons poitht clustering. Unlike fixed
13

---

UWt: Optimizing Ultrametric Trees for Tree-Wassesstiin Distance
IV. Tree-Sliced Methods: Inspired by the sliced- Wasserstein distance (Rabin al., 2012), Le et al. (2019) introduced the tree-sliced Wasserstein distance, which in- cludes the Sliced-Quad T'ree and Sliced-ClusterTre meth-
A.5. Other Related Work
Beyond the comparison methods, trees and hierarchical partitions have been widely explored for approximating the 1-Wasserstein distance. These include an importance sampling method (Indyk, 2007), greedy tree approaches (Khesin et al., 2021; Agarwal et al., 2023; Fox & Lu, 2023; Agarwal et al.. 2024), and a streaming algorithm (Chen et al., 2022). However, these works primarily rely on quadtrees, whereas our method focuses on ultrametric trees.
ods. These methods compute the average tree- See-Wassenstein distance (TWD) over multiple randomly constructed trees. Similarly, Yamada et al. (2022) proposed sliced versions of qTWD and cTWD, referred as Sliced-qIWD and Slieed-eTWD, respectively. The sliced TWD is defined as:
K
1 WT(n,)) = K
Σ WT, (1,
(16)
In addition, there are several tree- Wasserstein distances that are not included in our comparison due to different goals and formulations: (1) Lin et al. (2025) learns meaningful tree- Wasserstein distance that captures latent feature hierar- chies by using hyperbolic diffusion LCA relations. (2) STW (Takezawa et al., 2021) is a supervised method that uses con- trastive loss to learn task-specific distances. Since it does not aim to approximate W1, it yields higher errors; for exam- ple, on the BBCSport dataset, STW produces a large RE-W of 0,643 and retrieval precision of 0.335. (3) Tran et al. (2024) and Tran et al. (2025) extend the wereed-Wasersteein distance (Rabin et al., 2012) by using structured systems of lines that can be metrized by tree distances. While inter- pretable as tree-based -esseeseesn variants, they are not typical TWD methods for approximating W1-
where Wr, (p, 2)) represents the TWD computed on the i-th tree, and K is the total number of constructed trees.
V. Supervised Method: Ultra (Chen et al., 2024) in- troduced a supervised approach based on ultrametric trees, where the ultrametric matrix Dr = (d zy)] E R"xn is optimized by solving a regression problem:
min DrER"xn
(W ( MG)) Wy (مn, ) (17) (LJ)ES
Σ
where S denotes the training set, Dr represents the learned ultrametric, Witain is the 1-Wasserstein distance pre-computed on the training data, and Warsin denotes the tree-Wasserstein distance defined on the ultrametric tree T. As supervised approach, it requires extensive training data and the pre-computation of the 1-Wasserstein distance to serve as labels, which can be computationally expensive.
In contrast, UltraTWD is specifically designed to accurately approximate the I-Wasserstein distance by learning both tree structure and edge weights within a unified ultrametric framework.
Limitations of Existing Methods. The existing Wasser- stein approximation methods have several notable lim- itations, as summarized in Table 7: (1) The entropy- based method (Sinkhorn distance) deviates from the true I-Wasserstein distance due to the entropic regularization term and has a time complexity of O(n)). (2) Tree- construction methods cannot optimize the tree structure and edge weights and require the embedded data points X = [.,,.... xn. to partition the data space. (3) Weight- optimized methods also cannot optimize the tree structure and depend randomly constructed QuadTree or Clus- terTree. (4) Tree-sliced methods require multiple trees but still lack the capability to optimize tree structures. Addition- ally, these unsupervised TWD methods, based on QuadTree or ClusterTree, rely on embedded data points. (5) The su- pervised method can optimize tree structures and utilize the distance matrix D as input, eliminating the need for embed- ded data points. However, they require extensive training data and precomputed 1- Wasseerstei distances for training.
B. Experimental Setup
B.1. Datasets
We evaluate the performance using four benchmark text datasets (Huang et al., 2016): BBCSport, Reuters, Ohsumed, and Recipe. In these datasets, each word is represented as a 300-dimensional word2vec vector, while each document is represented by a normalized bag-of- distribution based on word frequencies. The details are provided below.
This dataset consists of 737 BBC sports articles labeled by classes. It contains five test sets with 220 articles each, averaging 6,051 words per test set.
BBCSport:
Reuters: This dataset contains 7,674 news articles across classes We randomly generate five test sets with 1,000 articles each, averaging 6,416 words per test set.
This dataset contains 9,152 medical abstracts within 10 classes We randomly create five test sets with 1,000 abstracts each, averaging 9,467 words per test set.
Ohsumed:
In contrast, the proposed UttraWWD methods inherit the advantages of UltraTree and can optimize both the tree structure and edge weights. Moreover, they surpass Ultra- Tree by eliminating the need for training data while offering greater accuracy and efficiency.
Recipe: This dataset consists of 4,370 recipe procedures with 15 classes. It contains five test sets with 1,311 tweets each, averaging 4,084 words per test set.
14

---

UWD: Optimizing Ultrametric Trees for Pree-Wasscrstein Distance
QuadTree
Tree
UltraTree UltraTWD
QuadTree
ClusterTree UltraTree UltraTWD
0.9
0211
Best 863
Be57=0.847 B85
0.9
Rest=0.824
Best Best 0.839 8.56=0.876.
Precision
0.8
Precisio
0.7
0.8
0.7
0.6
0.6
0.5
Sinktarn
QuadTre Silied Quad
Slicdd-7TV
gTWD
Cluster free Siceed Claster free
<TWD Siiced cTWD
UltraTree 1% Ultra 2%
u ultratree 10%
ulira Ultra wraTWD-MM MST litraTWD.o GD UltraTWD-I
Sinkhorn
QuadTree Siiced QuadTree
Sliced-ap ATWD
qTWD
Sliced Sliced-Clustrpi Clustertres cTWD
Clustertree
Silced cWDD
UltraTee 1%
Ultra Tree- 2% Ulira Traar 10ae-100%
UltraWD MST
UltraTwD GD
Ultra
(a) BBCSport
(b) Reuters
QuadTree 0.5t=0.708
ClusterTree UltraTree UltraTWD 0.52=0.742 8e5t=0.752 D85
QuadTree
ClusterTree UltraTree UitraTWD Bes1=0.a33 Best=0.866.
0.8
0.9
Best=0.801
Best=
Precision
0.7
Precision
0.7
0.6
0.5
0.5
0.4
0.3
0.3
0.1
Sinktorn
dTWD Sliced Quad
Tued
Sliced TWW
Slcadd Cuater Trwe
cTWD Silced-cTWW CWD
Ultra 1% UttraTree free-224 Ultratree 10%
UlfraTWD GO
MST
UltraTWD-ip Ip
Sinkham
Siiced Coungaee qTWD Sliced-q 4TWD
Sliced Cluster CTWD
Clstertiss
Silced-cTVe TWD
Ultra Ultratree 10%
1%
UltaaWo MST
UltraTWD-c GD UltraTWD-i
Ultra Tree
Ultra
(c) Ohsumed
(d) Recipe
Figure 6. Document retrieval precision of tree- Deesssssrrnt distance methods. Blue represents QuadTree-based methods, green represents
ClusterTree-hased methods, yellow represents UltraTree methods, and represents UltraTWD methods, Our UltrTTW-II GD approaches consistently outperform all baseline methods across four datascts.
B.2. Implementation and Hyperparameters
of 32, the Adam optimizer with learning rate of 0.01, and maximum of iterations (typically converging within 2).
The hyperparameters of baseline methods are listed below:
Sinkhorn method: The Sinkhorn distance is computed us- ing the ot.. i rn2 function from the Python Optimal Transport library (Flamary et al., 2021), with a regulariza- tion parameter of À = in Eq. (14) and maximum number of iterations numI e rmax = 100.
C. Comprehensive Results and Analysis
This section presents a detailed analysis of our comprehen- sive results, including
document retrieval performance (Section C.1).
hyperparameter and convergence analysis (Section C.2).
We implement QuadTree, ClusterTree, TTWD, cTWD, and their sliced versions using the official c provided by Yamada et al. (2022). For qTWD and cTTD,, the regularization parameter in Eq. (15) is set to À = 0.001, following the configuration in Yamada et al. (2022). For the tree-sliced methods, the number of multiple trees is set to K in Eq. (16), consistent with Yamada al. (2022).
Unsupervised TWD methods:
C.1. Document Retrieval Performance
We evaluate the average precision for the top-10 retrieved candidates and present the numerical results in Table 3(b). To provide deeper insights, the corresponding visualizations are shown in Fig. 6.
The key observations are as follows:
Supervised TWD method: For UltraTree, five training sets are randomly generated for each dataset, each containing 1,000 document distribu- tions. Each document distribution is created using numpy. random. randint nt (1, 11, size= (1, wowrrd) with 1% the entries randomly selected to form a sparse vector, which is then normalized to probability distribution. The official code is used with default settings: batch size
Effect Optimized Edge Weights: The qTWD and cTWD methods optimize edge weights based on the QuadTree and ClusterTree structures and thus significantly enhance retrieval performance compared to their non- optimized counterparts (QuadTree & ClusterTree).
The sliced methods improve precision by averaging tree- See-Wasserstein distances across multiple randomly constructed trees. However, without slicing, the retrieval performance of QuadTree is highly sensitive to randomness, leading greater variance.
Impact of Sliced Distances:
'nttps: :ttttps:/github. com/oist/tree07 "https://on /ttp://gtttttt thub com/ wwwnch/tneeeeeeeeee
https///trre___eeanneeeee Learning
15

---

UltraTWD: Optimizing Ultrametric Trees for Tree-Wassersterr Distance
While the sparsity of training data influences the approximation error of Wr, UrrraTree consistently maintains stable retrieval performance across varying sparsity levels (e.g., 1 1%, 2%. 10%).
Stability of UltraTree::
As shown in Fig. 7, both UttaTWDD-IP and GD methods exhibit empirical convergence across all four datasets. Re- gardless of whether iterative projection gradient descent is employed, the objective function || Dr - D[] is consis- tently minimized to comparable values. Although global optimality is challenging to guarantee, the Dr. obtained us- ing these methods yields high-quality Wr, resulting in high retrieval precision. The precision follows the order of IP > GD (a 0.02) MST > most baselines, validating the effectiveness of our methods.
UltraTWD Superiority: Among Uhr methods, UltraTWD-IP and GD consistently outperform UttraTW-- MST and all baseline methods. Compared to the best- performing baseline, UltraWW-IIP achieves significant pre- cision gains of 0.022, 0.027, 0.036, and 0.033 on the BBC- Sport, Reuters, Ohsumed, and Recipe datasets, respectively.
MST - w GD (a=0.01)
GD (a=0.02)
GD (a=0.05)
The improved precision validates that our Ult meth- ods effectively preserve neighbor relationships, ensuring that documents similar in W1 remain similar in WT.
0.3 BBCSport
BBCSport
0.89
Precision
0.88
0.87
C.2. Hyperparameter and Convergence Analysis
0.1
0.86
Hyperparameter Analysis. We analyze the impact of hy- perparameters on convergence and retrieval performance. (1) UtraTWD-MST: It directly computes the optimal solu- tion using Theorem 5 and requires no hyperparameters. (2) Ultra TWD-IP: The iterative projection method in Algo- rithm 2 has only one hyperparameter: the maximum number of iterations m. As shown in Fig. 7, retrieval precision slightly improves with more iterations on BBCSport, re- mains stable on Reuters, and decreases slightly on Ohsumed and Recipe datasets. Qualitatively, decline in performance with additional iterations may indicate overfitting of the Dr solution, weakening Wr's 's ability to preserve neighbor relationships. Despite these fluctuations, UtrrTTWD-IP con- sistently achieves the highest retrieval precision, outperform- ing TWD-GD and all baseline methods. Considering the time cost of iterative projection, we set m = and use results from a single iteration in practical applications, as already delivers excellent performance.
0.0
0.85
5
10
15
20
5
10
20
Iteration
Iteration t
Reuters
Reuters
0.3
0.89
Precision
0.88
0.2
0.87
0.1
0.86
0.85
0.0
5
10
15
20
5
10
15
20
Iteration
Iteration t
Ohsumed
Ohsumed
0.3
0.80
Precision
0.79
RE-D
0.2
0.78
0.77
0.0
0.76
5
10
15
20
5
10
15
20
Iteration t
Iteration
t
Recipe
Recipe
(3) TWD-G: The gradient descent method in Algo- rithm 3 has two hyperparameters: the learning rate and the maximum number of iterations m. As shown in Fig. 7, = 0.02 achieves faster convergence than a = 0.01 and smoother convergence than a = 0.05, leading to higher precision than a = 0.01 and more stable performance than 0.05. A large learning rate (a 0.05) may cause overshooting or oscillations around local minima, leading to suboptimal solutions. Empirically, the choice of depends on the specific dataset, and we use = 0.02 in this paper. For the maximum number of iterations, we set m 8 in practice, balancing low RE-D and high precision.
0.3
0.87
Precision
0.86
0.2
0.85
0.1
0.84
0.0
0.83
0
5
10
15
20
0
5
10
15
20
Iteration
Iteration
Figure 7. Hyperparameter analysis of UWt methods across four datasets, illustrating empirical convergence.
Convergence Analysis. To illustrate the optimization pro- cess, we evaluate the relative error of the learned ultrametric Dr from different iterations, defined as
Dr DlEF
RE-D-
||D
which is proportional to the objective function Dr D/
16


Reviews:
[{"review_id": "FJNvdqwwSt", "reviewer": "Reviewer_SjiH", "summary": "The paper proposed a method to find out an Ultra tree Wasserstein distance. The method is based on minimizing a distance between trees satisfying certain conditions for trees. Algorithm 2 proposed to find the solution under projections, meanwhile Algorithm 3 tries to reduce the computation when avoiding working with pairwise distance between leafs, instead it works with the node heights. The second last section is devoted for testing the proposed method on several datasets.\n\n##update after rebuttal: I keep my score, since the authors mostly explained the difficulties of the problem.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your valuable comments. We will modify it accordingly.\n\n---\n\n**Comment 1.** *No guarantee for the convergence of the proposed algorithms.*\n\n**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.\n\n[1] Chierchia, G. and Perret, B. Ultrametric fitting by gradient descent. NeurIPS, 2019.\n\n[2] Chen, S., Tabaghi, P., and Wang, Y. Learning ultrametric trees for optimal transport regression. AAAI, 2024.\n\n---\n\n**Comment 2.** *Algorithm 2 would be costly because the number of projections is too large.*\n\n**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).\n\n---\n\n**Comment 3.** *For the algorithm 3, theoretically is the $D_T$ determined uniquely by $H_T$? Since the number of entries reduce significantly?*\n\n**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.\n\n---\n\n**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.*\n\n**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.**\n\n---\n\n**Comment 5.** *UltraTree also produce quite competitive results, does this mean that the proposed method not really improve the overall performance significantly enough?*\n\n**Response 5.** While UltraTree is a strong supervised baseline, our UltraTWD methods show clear advantages in performance, generalizability, and efficiency.\n\n- **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).\n\n  **Table R1. Precision comparison of document retrieval.**\n  |Dataset|BBCSport|Reuters|Ohsumed|Recipe|\n  -|-|-|-|-\n  Best Unsupervised|0.863|0.849|0.742|0.831\n  UltraTree|0.842|0.834|0.749|0.830\n  UltraTWD-GD|0.868|0.860|0.776|0.848\n  UltraTWD-IP|**0.885**|**0.876**|**0.788**|**0.866**\n  Improvement $\\uparrow$|+5.1%|+5.0%|+5.2%|+4.3%\n\n- **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).\n\n- **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.\n\n**In summary, UltraTWD methods are more accurate, efficient, and broadly applicable than UltraTree.**\n\n---\n\n**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?*\n\n**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).", "claims_and_evidence": "New framework for tree-Wasserstein distance exploit the ultra metric property. \n\nThe problem is formulated into ultra metric optimization  questions and proposed algorithm to solve it\n\nEmpirical evidences are shown to demonstrate the effectiveness of the methods. \n\nNo guarantee for the convergence of the proposed algorithms.", "methods_and_evaluation_criteria": "The proposed methods sound reasonable. The metrics are fine.", "theoretical_claims": "All theoretical results are standard in theory of tree Wasserstein distance. The authors do not propose any new theoretical result.  \nAlgorithm 2 would be costly because the number of projections is too large.\n\nFor the algorithm 3, theoretically is the $D_T$ determined uniquely by $H_T$? Since the number of entries reduce significantly?", "experimental_designs_or_analyses": "According to Table 2, the UltraTWD-IP and UltraTWD-GP 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. In the same table, UltraTree also produce quite competitive results, does this mean that the proposed method not really improve the overall performance significantly enough?", "relation_to_broader_scientific_literature": "Yeap, I think it is related to all other topics involving tree-data.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "No", "other_comments_or_suggestions": "Why in Algorithm 2, line 3: the weights are set to be $\\frac{1}{t+1}$ and $\\frac{t}{t+1}$ in which $t$ is the $t$th iteration?", "questions_for_authors": "Please read the above comments", "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": "MlYImf3p9P", "reviewer": "Reviewer_dNo2", "summary": "1. Wasserstein distance has been applied to many tasks. This paper mainly focuses on how to learn an optimal tree-Wasserstein distance, which is not limited to a specific task.\n2. The primary motivation of this paper is to address the suboptimal tree structures and inadequately tuned edge weights in traditional tree-Wasserstein distance. The proposed new framework (UltraTWD) is the first unsupervised framework to simultaneously optimize both tree structure and edge weights by leveraging the ultrametric property.\n3. The authors formulate ultrametric nearness problems to optimize trees equipped with the nearest ultrametric to a cost matrix and propose efficient algorithms to address them. These algorithms are based on minimum spanning trees, gradient descent, and an accurate method using iterative projection, all of which deliver high-quality solutions.\n4. The proposed new framework achieves the lowest estimation errors compared to both unsupervised methods and the state-of-the-art\nsupervised method across four benchmark datasets. Additionally, it demonstrates exceptional performance in document retrieval, ranking, and classification, showcasing its practicality for Wasserstein-distance-based applications", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n\n**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?*\n\n**Response 1.** Thank you for raising this important point. we apologize for any confusion and would like to clarify:\n\nIn 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.\n\n**Table R1. Comparison of document classification accuracy.**\n|Dataset|BBCSport|Reuters|Ohsumed|Recipe|\n-|-|-|-|-\nBest baseline|0.826|0.902|0.423|0.495\nOur distance|0.839|0.905|0.428|0.495\nTradition distance|**0.874**|**0.919**|**0.498**|**0.499**\n\n---\n\n**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.*\n\n**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...*\"\n\nDespite 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.\n\n[1] Chierchia, G. and Perret, B. Ultrametric fitting by gradient descent. NeurIPS, 2019.\n\n[2] Chen, S., Tabaghi, P., and Wang, Y. Learning ultrametric trees for optimal transport regression. AAAI, 2024.\n\n---\n\n**Comment 3.** *How close the proposed distance is to the traditional distance has not been analyzed theoretically.*\n\n**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:\n\\begin{equation}\n|W_T(\\mu, \\nu) - W_1(\\mu, \\nu)| \\le 2||D||\\_\\infty + ||D_T - D||\\_\\infty \\le 2||D||\\_\\infty + ||D_T - D||\\_F\n\\end{equation}\nThis 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.\n\n---\n\n**Comment 4.** *The generalization analysis of the proposed new method is also lacking.*\n\n**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).\n\nAs 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$.\n\n**Table R2. Approximation error (RE-W) under different test sparsity.**\n|# valid words ($n$)|1000|2000|3000|4000|5000|\n-|-|-|-|-|-\nSparsity|17%|33%|50%|66%|83%\nBest baseline|0.116|0.150|0.167|0.176|0.176\nUltraTWD-GD|0.088|0.111|0.122|0.127|0.125|\nUltraTWD-IP|**0.065**|**0.073**|**0.076**|**0.076**|**0.072**", "claims_and_evidence": "Yes. The claims in the submission are supported by clear and convincing evidence, such as the complexity analysis and experiments.", "methods_and_evaluation_criteria": "I think that the proposed methods and/or evaluation criteria (e.g., benchmark datasets) make sense for the problem or application at hand.", "theoretical_claims": "All the author's theoretical analysis comes from other references, so I think there is no need to check the correctness of any proofs for theoretical claims.", "experimental_designs_or_analyses": "The author's experiment is relatively simple, and I don't think there are any significant problems in these experiments.", "relation_to_broader_scientific_literature": "Compared with other related literature, the key contribution of this paper is to improve the accuracy of traditional methods. I think this paper has strong application and may be applied to other fields of natural science, such as image analysis, pattern recognition, etc.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths: \n1. The motivation of this paper is clear. The authors mainly address the limitations of suboptimal tree structures and inadequately tuned edge weights in traditional methods. The studied problem is meaningful. \n2. This paper is well organized, logically clear, and written fluently, making it easier for readers to understand.\n3. In this paper, the authors address a key problem: How to bridge the gap between WT and W1. A series of algorithms have been proposed to solve this problem. The author calculated the complexity of the method and showed that his method has lower complexity. Combining minimum spanning trees, iterative projection, and gradient descent provides a robust and innovative approach to optimizing tree structures and edge weights.\n4. The current experiments are Ok.\n\nWeaknesses:\n1. The rationale for the core idea is not fully explained. The authors formulated the tree-metric nearness problem to closely approximate D.  Specifically, the authors use Wasserstein distance as a constraint to control this tree-metric. This idea is good. The authors aim to make the TWD distance as close as possible to a certain distance to improve its accuracy. There is a problem here:  The authors use approximation strategies throughout the paper but do not analyze the difference between this approximation and real Wassestein distance. 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?\n2. Theoretical analysis is not enough. 1）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. A more rigorous theoretical analysis would enhance the credibility of the framework. 2）How close the proposed distance is to the traditional distance has not been analyzed theoretically. For example, there is no theoretical analysis of the error between your new method and the Wasserstein distance. 3）The generalization analysis of the proposed new method is also lacking.", "other_comments_or_suggestions": "No", "questions_for_authors": "See Weaknesses", "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": "BGaOeBA1VS", "reviewer": "Reviewer_XRdY", "summary": "The Wasserstein distance is a well-known metric for comparing distributions, and has been used as loss function in many ML models. To improve the efficiency of computing the Wasserstein distance, researchers have considered embedding the distributions to a tree metric and computing the Wasserstein distance over the tree embedding, which can be done in linear time. This paper proposes to embed the distributions into a tree that has ultrametric distances, and provide three algorithm for computing the ultrametric tree distances that are close to the underlying metric.\n\nMore specifically, the authors formulate the nearness of a tree metric to the ground metric in two ways, namely the $\\ell_\\infty$ norm and the Frobenius norm of the tree distance and the ground distance. For the $\\ell_\\infty$ norm, they present a simple algorithm that computes the tree ultrametric using minimum spanning trees. For the Frobenius norm, they present an IPM-based algorithm and a gradient-descent-based approach to compute the ultrametric. \n\nThey show in experiments on text datasets that their proposed tree Wasserstein distance (Ultra-TWD) has a lower approximation ratio as well as better performance in document retrieval, ranking, and classification tasks compared to other variants of the tree Wasserstein distance in the literature.\n\n\n## Update after rebuttal\nThe authors resolved my concerns. They provided additional experimental results and promised to add the new experimental results as well as the missing citations. I increased my score from 2 to 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We truly appreciate your positive feedback. Your valuable suggestions will help us improve the work.\n\n---\n\n**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.*\n\n**Response 1.** To ensure a fair comparison, we consider the total time complexity required to compute $W \\in \\mathbb{R}^{N \\times N}$.\n- For exact $W_1$: $O(N^2 \\cdot n^3 \\log n)$.\n- For UltraTWD: (including the tree learning time)\n  - $O(N^2 \\cdot n + n^3)$ for UltraTWD-IP\n  - $O(N^2 \\cdot n + n^2)$ for UltraTWD-GD\n\nAs 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:\n- UltraTWD-GD learns the tree in 34s and computes each $W_T(\\mu, \\nu)$ in 0.4s using $O(n)$ time,\n- Each $W_1(\\mu, \\nu)$ takes 1.7s with $O(n^3 \\log n)$ time.\n\nTo 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.\n\n---\n\n**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.*\n\n**Response 2.** We would like to clarify:\n- 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.\n- 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.\n\nAs 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.\n\n**Table R2. Performance comparison on the BBCSport dataset.**\n|Metric|RE-W$\\downarrow$|Precision$\\uparrow$|MRR$\\uparrow$|ACC$\\uparrow$|\n-|-|-|-|-\nSliced_qTWD (3 trees)\n$\\lambda=0.01$|0.142|0.809|0.859|0.821\n$\\lambda=0.1$|0.118|0.798|0.852|0.815\nSliced_cTWD (3 trees)\n$\\lambda=0.01$|0.108|0.867|0.898|0.820\n$\\lambda=0.1$|0.036|0.823|0.870|0.832\nUltraTWD-GD|0.016|0.868|0.921|0.838\nUltraTWD-IP|**0.014**|**0.885**|**0.924**|**0.839**\n\n[1] Chen, S., Tabaghi, P., and Wang, Y. Learning ultrametric trees for optimal transport regression. AAAI, 2024.\n\n[2] Yamada, M., Takezawa, Y., Sato, R., Bao, H., Kozareva, Z., & Ravi, S. Approximating 1-wasserstein distance with trees. TMLR, 2022.\n\n---\n\n**Comment 3.** *I believe the following papers can be cited in this work.*\n\n**Response 3.** We will cite these relevant literature in the revision.\n\n---\n\n**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.*\n\n**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.\n\n---\n\n**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).*\n\n**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**.\n\n**Table R3. Performance comparison on the BBCSport dataset.**\n|Metric|RE-W$\\downarrow$|Precision$\\uparrow$|MRR$\\uparrow$|ACC$\\uparrow$|\n-|-|-|-|-\nSOPT|113.6|0.082|0.079|0.357\nSPOT|0.666|0.064|0.063|0.231\nUltraTWD-GD|0.016|0.868|0.921|0.838\nUltraTWD-IP|**0.014**|**0.885**|**0.924**|**0.839**\n\n[1] Bai, Y., Schmitzer, B., Thorpe, M., & Kolouri, S. Sliced optimal partial transport. CVPR, 2023.\n\n[2] Bonneel, N., & Coeurjolly, D. Spot: sliced partial optimal transport. ACM TOG, 2019.", "claims_and_evidence": "Yes, most claims are well-supported. I have raised some concerns about theoretical and experimental results below.", "methods_and_evaluation_criteria": "Yes, the datasets are suitable for the problem.", "theoretical_claims": "I have a concern with one of the claims made in Section 3.6: The UltraTWD (or any other tree Wasserstein distances) provides an approximation of the Wasserstein distance. Therefore, comparing the running time of $W_T$ with the exact computation of $W_1$ is unfair. The UltraTWD and an approximation algorithm such as Sinkhorn would have roughly the same execution time (although it is not clear how many iterations does the GD algorithm require).", "experimental_designs_or_analyses": "I have a concern with the experiments. Specifically, I am concerned that the experimental setup for competing distances might not be set optimally. For instance, 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, as I expect to see better performance by setting $\\lambda=0.01$ and max number of iterations to 300 for instance. 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.", "relation_to_broader_scientific_literature": "The problem of computing Wasserstein distance is a well-studied area, and designing fast algorithms that approximate the Wasserstein distnace is of high importance. Approximating the ground distances using a tree metric can help speeding-up the computation of Wasserstein distance, as the Wasserstein distance on a tree metric has a simple closed-form formula that can be computed in linear time.", "essential_references_not_discussed": "The use of trees and hierarchical partitionings in approximating the 1-Wasserstein distance has been extensively used in the literature, and I believe the following papers can be cited in this work:\n\nTree-based Algorithms:\n* P. Indyk. \"A near linear time constant factor approximation for Euclidean bichromatic matching (cost).\" SODA 2007.\n* P. K. Agarwal, S. Raghvendra, P. Shirzadian, and R. Sowle. \"A higher precision algorithm for computing the 1-Wasserstein distance.\" ICLR 2023.\n\nMWU-based approaches for boosting the accuracy of greedy tree algorithms:\n* A. Khesin, A. Nikolov, and D. Paramonov. \"Preconditioning for the Geometric Transportation Problem.\" SOCG 2019.\n* E. Fox and J. Lu. \"A deterministic near-linear time approximation scheme for geometric transportation.\" FOCS 2023.\n* P. K. Agarwal, S. Raghvendra, P. Shirzadian, and K. Yao. \"Fast and accurate approximations of the optimal transport in semi-discrete and discrete settings.\" SODA 2024.\n\nStreaming Algorithm:\n* X. Chen, R. Jayaram, A. Levi, and E. Waingarten. \"New streaming algorithms for high dimensional EMD and MST.\" STOC 2022.", "other_strengths_and_weaknesses": "**Strengths**\n* The paper is well-written and easy to follow. \n\n**Weaknesses**\n* There are no theoretical guarantees on the convergence rate of their proposed IPM and SGD algorithms, and hence, it is hard to judge whether the proposed methods have a comparable time as the previous methods or not. \n* 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).", "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": "cKTXiV2i8I", "reviewer": "Reviewer_XieA", "summary": "This paper introduces an unsupervised approach to constructing an ultrametric tree for the Tree-Wasserstein Distance (TWD) that approximates the Wasserstein distance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable questions. **We will revise the paper accordingly and add more discussion on related work.**\n\n---\n\n**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?*\n\n**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.\n\n[1] Chen, S., Tabaghi, P., and Wang, Y. Learning ultrametric trees for optimal transport regression. AAAI, 2024.\n\n---\n\n**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?*\n\n**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.\n\n---\n\n**Question 3.** *Could you provide a visualization of the constructed tree from the proposed methods?*\n\n**Answer 3.** Please see the [figure](https://anonymous.4open.science/r/rebuttal_twd/visualization.pdf). We will include it in the revision.\n\n---\n\n**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?*\n\n**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:\n- **$O(N^2 \\cdot n + n^3)$** for UltraTWD-IP\n- **$O(N^2 \\cdot n + n^2)$** for UltraTWD-GD\n- **$O(N^2 \\cdot n^3 \\log n)$** for traditional $W_1$. \n\nFor 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.\n\n---\n\n**Question 5.** *When reporting the computation time in Figures 7 and 8, which algorithm was referred to in UltraTWD?*\n\n**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.\n\n---\n\n**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?*\n\n**Answer 6.** As suggested, we will include the following comparisons:\n\n- **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.\n\n  **Table R1. Performance comparison on the BBCSport dataset.**\n  |Metric|RE-W$\\downarrow$|Precision$\\uparrow$|Total Time (hour)|\n  -|-|-|-\n  STW|0.643|0.335|3.5\n  UltraTree|0.022|0.842|0.6\n  UltraTWD-GD|0.016|0.868|**0.1**\n  UltraTWD-IP|**0.014**|**0.885**|0.5\n\n  [1] Takezawa, Y., Sato, R., & Yamada, M. Supervised tree-wasserstein distance. ICML, 2021.\n\n  [2] Chen, S., Tabaghi, P., and Wang, Y. Learning ultrametric trees for optimal transport regression. AAAI, 2024.\n\n- **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.\n\n  **Table R2. Comparison of SVM classification accuracy.**\n  |Dataset|BBCSport|Reuters\n  -|-|-\n  WMD|**0.874**|**0.919**\n  UltraTWD-GD|0.838|0.900\n  UltraTWD-IP|0.839|0.905", "claims_and_evidence": "- The method relies on a tree-based representation, yet the paper does not clearly define what constitutes the tree (e.g., whether it is strictly binary or can be flexible) or how the structure is updated. \n- The optimization goal is related to metric multidimensional scaling, but there is no discussion about it.\n- Although the authors claim an efficient approximation of the Wasserstein distance, the complexity analysis (especially the O(n³) cost for UltraTWD-IP) does not clearly indicate a substantial improvement over the O(n³ log n) complexity of traditional optimal transport methods.", "methods_and_evaluation_criteria": "The proposed methods appear generally appropriate for document data applications and the use of Wasserstein distance approximations. The experimental evaluation is a strength; however, since the paper posits the paper’s competing work with the supervised method, it would be better to also include the results from the supervised tree-Wasserstein distance by Takezawa et al. and discuss this related work. In addition, since the paper claims to be an efficient approximation of Wasserstein distance and the application is document data, the empirical results of WMD should be included in the comparison.", "theoretical_claims": "The complexity claim for UltraTWD-IP is stated as O(n³). The improvement over classical optimal transport (O(n³ log n)) is not significant enough to be a clear theoretical advantage.", "experimental_designs_or_analyses": "- Document datasets are appropriate, and the empirical evaluation is comprehensive regarding performance metrics.\n- The absence of experiments comparing with supervised tree-Wasserstein distance and WMD limits our understanding of the method’s benefits. \n- There is no clear explanation of how the tree structure is updated.", "relation_to_broader_scientific_literature": "The paper contributes to the literature on efficient approximations of the Wasserstein distance for document data.", "essential_references_not_discussed": "- The triplet constraints or triplet relation were explored in previous work, e.g., An Improved Cost Function for Hierarchical Cluster Trees. Although they have not yet been applied within iterative projection frameworks, it would be valuable to reference this earlier research considering the triplet in tree relation. \n- Related work not discussed:\n    - Distance-Based Tree-Sliced Wasserstein Distance\n    - Tree-Wasserstein Distance for High Dimensional Data with a Latent Feature Hierarchy\n    - Projection Optimal Transport on Tree-Ordered Lines", "other_strengths_and_weaknesses": "Strengths:\n- The paper is well-written and easy to follow.\n- Extensive experimental results on document datasets.\n\nWeaknesses:\n- Minor issues include undefined symbols (e.g., $n$ in Section 1, the norm in Eq. (7), and the height in Eq. (8)).", "other_comments_or_suggestions": "NA", "questions_for_authors": "- 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?\n- How is the tree structure updated during the iterative projection process? Does the update involve regrouping leaves or modifying the overall tree topology?\n- Can you elaborate on the practical benefits of the O(n³) complexity for UltraTWD-IP, especially when compared to the O(n³ log n) complexity of traditional optimal transport methods?\n- When reporting the computation time in Figures 7 and 8 in Appendix C.3, which algorithm was referred to in UltraTWD?\n- Could you include empirical comparisons with supervised tree-Wasserstein distance methods and standard WMD, given that your method claims to be an efficient approximation?\n- Could you provide a visualization of the constructed tree from the proposed methods?", "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:

Meta-review: This paper proposes a new approximation algorithm for the tree-Wasserstein distance. More specifically, the paper proposes a computationally efficient algorithm to approximate a distance by tree-metric and use it for tree-Wasserstein distance. The key difference from the existing method is that the proposed method optimizes both the tree structure and its edge weight. Through experiments, it demonstrates that the proposed approximation method outperforms existing approaches.

The proposed approach is interesting and practically useful. Some reviewers have pointed out that the paper lacks novel theoretical results. However, the proposed method is highly useful, and its effectiveness is well validated by experiments. Most reviewers support the paper, and I also support its acceptance.
Decision: Accept (poster)