Tree-Sliced Entropy Partial Transport
Authors: Viet-Hoang Tran
Thanh Tran
Thanh Chu
Tam Le
Tan Minh Nguyen
Venue: NeurIPS 2025 poster ()
Keywords: tree-sliced wasserstein distance
optimal transport
partial transport
unbalanced optimal transport
Primary area: general_machine_learning
TL;DR: We propose a tree-sliced framework to the Partial Transport setting by observing that the Partial Transport problem on tree-metric spaces can be reformulated as a standard Optimal Transport problem.

Abstract: Optimal Transport (OT) has emerged as a fundamental tool in machine learning for comparing probability distributions in a geometrically meaningful manner. However, a key limitation of classical OT is its requirement that the source and target distributions have equal total mass, limiting its use in real-world settings involving imbalanced data, noise, outliers, or structural inconsistencies. Partial Transport (PT) addresses this limitation by allowing only a fraction of the mass to be transported, offering greater flexibility and robustness. Nonetheless, similar to OT, PT remains computationally expensive, as it typically involves solving large-scale linear programs—especially in high-dimensional spaces. To alleviate this computational burden, several emerging works have introduced the Tree-Sliced Wasserstein (TSW) distance, which projects distributions onto tree-metric spaces where OT problems admit closed-form solutions. Building on this line of research, we propose a novel framework that extends the tree-sliced approach to the PT setting, introducing the Partial Tree-Sliced Wasserstein (PartialTSW) distance. Our method is based on the key observation that, within tree-metric space, the PT problem can be equivalently reformulated as a standard balanced OT problem between suitably modified measures. This reformulation enables efficient computation while preserving the adaptability and robustness of partial transport. Our method proves effective across challenging tasks such as outlier removal and addressing class imbalance in image-to-image translation. Our code is publicly available at [this https URL](https://github.com/thanhqt2002/PartialTSW).

Tree-Sliced Entropy Partial Transport
Viet-Hoang Tran' Department Mathematics National University of Singapore hoang. tranviettu. nus edu
Thanh Tran College of Engineering & Computer Science
VinUniversity 11tahhh. tq@vinuni edu.. vn
Thanh Chu School of Computing National University of Singapore thanh. chu@u, edu
Tam Le
Department of Advanced Data Science Institute Statistical Mathematics tam sm. ac jp
Tan M. Nguyen'
Department of Mathematics National University of Singapore tanmn@nus. edu. sg
Abstract
Optimal Transport (OT) has emerged as fundamental tool in machine learning for comparing probability distributions a geometrically meaningful manner. However, key limitation of classical OT is its requirement that the source and target distributions have equal total mass, limiting its use in real-world settings involving imbalanced data, noise, outliers, structural inconsistencies. Partial Transport (PT) addresses this limitation by allowing only a fraction of the mass be transported, offering greater flexibility and robustness. Nonetheless, similar to ОT, PT remains computationally expensive, as it typically involves solving large-scale linear programs-cve especialy high-dimensional spaces. To alleviate this computational burden, several emerging works have introduced the Tree- Sliced Wasserstein (TSW) distance, which projects distributions onto tree-metric spaces where OT problems admit closed-form solutions. Building on this line of research, we propose novel framework that extends the tree-sliced approach to the PT setting, introducing the Partial Tree-Sliced Wasserstein (PartialTSSW) distance. Our method based on the key observation that, within tree-metric space, the PT problem can be equivalently reformulated as standard balanced OT problem between suitably modified measures. This reformulation enables efficient computation while preserving the adaptability and robustness of partial transport. Our method proves effective across challenging tasks such as outlier removal and addressing class imbalance in maage-to-imaae translation. Our code is publicly available at https: //gi thub. com/
Introduction
Optimal Transport (OT) [87, 61] is framework for comparing probability distributions by lifting ground cost defined between individual points to a metric over measures. Its ability to capture the
"Corresponding author CCo-last author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

geometric structure of distributions leads widespread adoption across numerous domains, includ- ing machine learning [56. 11. 23, 32], data valuation [35, 38], multimodal data analysis [59, 49], statistics [50, 54, 57, 63], and computer vision [55, 70, 76, 89]. Despite its theoretical appeal, OT exhibits two major limitations in applications. First, the computational cost of OT problems for discrete measures scales as On 1og m) [61]. Second, the framework imposes a strict mass equality constraint, which is often violated due to noise, outliers, or unbalanced distributions [72].
Sliced Optimal Transport. To mitigate the high computational cost associated with OT, the Sliced Wasserstein (SW) distance [66, 10, 67] has been introduced as an efficient approximation. SW leverages the closed- foree-fommmae solution of one-dimensional OT by projecting high- sigh-dimensionl probability measures onto one me-dimensional subspaces, computing the OT cost in each slice, and subsequently averaging these costs. This procedure reduces the computational complexity to a sequence of sort- based operations with (n log n) complexity [61], while preserving key statistical and topological properties [52, 4, 29]. The SW framework has further inspired wide range of generalizations, including extensions based on structured projections [39, 18, 53), as well as adaptations to manifold- valued and non- Euclidean domains such the sphere [5, 65] and hyperbolic space [8].
Tree-Sliced Optimal Transport. One-dimensional projections, while computationally efficient, often fail to capture the intricate topological features inherent in high-dimensional data. To address this shortcoming, a growing body of work has explored richer integration domains as alternatives to linear projections in OT. These efforts span variety of metric settings, including Euclidean spaces [1, 60. 58], tree metries [46, 81], graph-based structures [45, 43], spherical geometries [65, 5, 83), and hyperbolic spaces [6, 48). A seminal contribution in this direction is the tree system proposed by [81], which serves as a structurally enriched substitute for traditional lines. By leveraging established results and closed-fom OT solutions on tree metric spaces [46, 34, 33], this framework introduces the Tree-Sliced Wasserstein (TSW) distance-a refined analogue of the classical SW distance. TSW retains the low computational complexity of SW while enhancing its capacity to reflect underlying data geometry. Recent advancements and extensions of the TSW framework are explored in [80. 79, 84).
Partial Transport and Unbalanced Optimal Transport. various applied settings. it often necessary to compare positive measures with unequal total for instance, in biological ap- plications where such measures represent cell populations of varying sizes [72]. The rigid mass conservation requirement of the classical OT can be relaxed using the Unbalanced OT (UOT) framework [73, 40, 47, 15, 22, 28], which introduces penalty terms that softly enforce mass preser- vation rather than enforcing it strictly. A related and widely used relaxation is Partial Transport (PT) [27, 9, 2, 68, 42, 43, 44], which allows only a fraction of the mass to be transported, thereby enabling more flexible alignment between distributions. PT improves robustness to outliers and facilitates meaningful comparisons under structural or statistical mismatches. It has also shown effec- tiveness in robust distributional alignment and has found applications in several domains, including deep learning theory [14, 69], cellular biology [72, 17], and domain adaptation [24, 3]. Despite its advantages, PT remains computationally intensive and is susceptible to noise in nigh-dimensional settings [20].
To address both the computational and mass imbalance issues, several approximate and scalable variants have been introduced, including entropic OT [16, 62] and minibatch OT [25, 24]. Recent work also has extended the sliced OT framework to the unbalanced setting [7, 22, 51, 28, 9, 2, 43], resulting in Sliced UOT variants with improved scalability and robustness.
Contributions. Motivated by the expanding TSW framework and recent advances in PT on tree metric spaces [42], this paper introduces a tree-sliced approach for computing partial transport between unbalanced measures in Euclidean spaces. The paper is organized as follows:
In Section 2, we review the foundations of Optimal Transport and Entropy Partial Transport metric spaces with tree metrics, well as the Tree-Sliced Wasserstein distance for probability measures in Euclidean spaces based on tree systems. These concepts collectively form the theoretical foundation upon which the proposed framework is developed
* In Section 3, we formally introduce the Tree-Sliced Entropy Partial Transport (PartialTSW) distance for comparing probability measures with unbalanced mass in Fuclidean spaces. We establish its metric properties and provide an analysis of its computational complexity.

---

In Section 4, we empirically evaluate PartialTSW on challenging tasks, such as enhancing noise robustness for generative models and addressing class imbalance in image-to- to-i translation. The results underscore its practical effectiveness and computational efficiency. We conclude our work in Section
All supplemental materials materials--inlluding theoretical foundations, formal proofs, experimental settings accompanied by additional tables and figures, and table of notation-are provided in the Appendix.
2 Building Blocks of Tree-Sliced Entropy Partial Transport
This section provides the foundations of Optimal and Entropy Partial Transport, as well as the Tree-Sliced Wasserstein distance. For the remainder of the paper, we denote the dimension by d.
2.1 Optimal and Entropy Partial Transports on Metric Spaces with Tree Metries
Tree Metric Spaces. Let = E) be tree rooted node r, with nonnegative edge lengths (weJees. We identify T with the set of all nodes and points along its edges. For a metric space Ω with metric is called free metric [75, 46] if there exists a tree 7 such that 92 T and d y) equals the length of the unique path between I and for all 2, E Q2 T is called a tree metric space. Assume V C Rd and let dr denote the tree metric on T; we write [z. y) for the unique path between y E T. Let w be the unique Borel (length) measure on T satisfying w( v) = dr(æ, y) for all I, T. For any I T, the subtree rooted is defined as A(x) = {y
[r. v... Figure I (left) provides a visual representation of tree metric spaces and their associated concepts.
Optimal Transport Tree Metric Spaces. Let P(T) be the collection of all probability measures on T (i.e. total mass is equal to 1). Let μ4, V/ E P(T), and P(u, v) be the set of coupling between μ and v. The L-Wasserstein distance (W) [87] between pl, is:
W p.dd (H, = inf (EEP(HMM) TxT dT y)2 dz(z.))
I
In the case of tree metrics and = 1, the distance W1,d7 (14, 1/) admits a closed-form solution [46]:
W1.d7 (1, v) - (I((A(=))=<AA(()))) - x(A()))
w(dx).
(2)
For general p > 1, the distance W Wp.d (14, v)) does not admit a closed-form expression. A natural generalization of Equation (2) leads the Sobolev Transport (ST) [45], that is
v(A(z)) STp(H, = (f/ - w(dx))
1 inf [ dr(x, dr =W Wpddy (H,
(3)
(EP((N)) TTX
Although ST, differs from Wp.rr it still defines a valid metric on P(T). Due to this property, and for simplicity, we focus on the case p 1, as the extension general p analogous.
Entropy Partial Transports on Tree Metric Spaces. Let M(T) and M(T 7) denote the space of all nonnegative Borel measures T and T x T respectively, with finite total mass. Given μ, E M(T). define the set of admissible partial couplings as
(u, = (7 M(T x T) 711 S4 72 v).
where ?1 and 7/2 represent the marginals of y on the first and second marginals, respectively. For any y E [I, v). let fi and f2 be the Radon Radon-Nikodil Nikodym derivatives of with respect to μ and 72 with respect v, respectively. Let w
[0, 2o) defined by w(x) = a da 2n) + ag where ao E T, di E b). and E [0. oo). We have is b-Lipschitz continuous. We use the entropy function F: [0, o0) (0, oo) given by F(s) Is 1.. Letting m
(((T and fixing
m E [0, m]. the Entropy Partial Transports (EPT) problem is formulated
frm Film (P) + Fa v) + TXT
Wm(uvv) inf 7Ells (n.e) (Tx7)==
di dr v) v)x,
(5)
3

---

82
T
Figure I: An illustration the construction of TSW. (Left) Tree Metric Space. The tree is rooted at node with all other nodes denoted by т;. Edges denoted by e₁. each assigned weight We. probability distribution on the tree assigns mass to its nodes. The subtree A(x) is defined the collection of all points lying along the edges in the subtree rooted at For example, A(r) includes the entire tree, (xx) includes all points on edges e7 and cg. and A(xg) includes all points edges e10, 611. and 612. (Right) An illustration of the Thee-Sliced Wasserstein computation. Given two probability measures (depicted in red and blue), and tree system T. each measure is first pushed-forward onto the tree via the Radon transform, resulting in measures supported the tree structure. The Wasserstein distance between these tree-projected measures is then computed using Equation (2). The overall TSW distance is obtained by averaging the Wasserstein distances across a collection of such trees, typically approximated using a Monte Carlo sampling framework.
where the regularization terms are defined as the weighted relative entropies
Fi(m p) = /, w(x) ((f((x)) p(dr), Fa(ma ()) / w(x) F(fa(z))
(6)
To handle the mass constraint 7( T) - m, a Lagrange multiplier introduced. Consider the relaxed objective
Fi(m (1) + F2 +)+ +6 fry, (dT(x, - X) 7(xx, dy)
(T) inf (ell
(7)
According to a construction by Caffarelli and McCann [12], the problem (7) is equivalent to a balanced OT problem. Utilizing duality and regularization techniques in [42], the objective (7) admits closed-form solution for the a-regularized EPT:
For E 60, + wrrr], it is given by:
(H) = 1 |m(A(=)) - v(A(x)) w(dx)
by ( (w(r) 6d
a) (u(T)-
(8)
Define the corresponding regularized transport cost
(a v) = = v) + by
2 (u(T) + v(]]
(9)
The function d, defines a metric on M(T), making (M(T), da) a complete metric space.
2.2 Tree-Sliced Wasserstein Distance on Euclidean Spaces
We adopt the setting of Tree-Sliced Wasserstein Distance on Systems of Lines as in [81, 80). Figure 1 (right) presents the illustration relevant the following discussion.
Tree System. A line in I is an element of Rd Sd-1, and system of k lines is an element of (RR x $d-1 We denote a system of lines by L, line in L (also used an index) by I, and the space of all such systems by L The ground set of L is defined
E (x XL: = to 9 for some t, E ER
where + + Oj with /. E R, is the parameterization of the line 1. For notational convenience, index the lines l,,..., .....,,,, lc, where each line defined by source point v; TR and direction vector 0, E gd-1, tree system is system of lines endowed with additional tree structure. To

---

highlight the presence of this structure, we denote the system by T rather than L. The space of tree systems-tth that is, the collection of tree systems sharing common tree structure- is denoted by TV or simply T. This space is equipped with probability distribution which is induced by a random sampling procedure over lines.
Radon Transform on Tree Systems. For T E TV denote the space of Lebesgue integrable functions on T as
L'(T) f: f:T- 1 E (f(tt.))) dt, < } 1ET
Define the space C(R. TV. 11-1) the set of continuous maps from Rd Ty the (k: 1)- dimensional standard simplex A named splitring maps. For L' (RR", define Rgf T - R such that:
R95(x11) = / f(v) alu: Th 51x - (y - x1, -1)) dy.
(10)
The function R is in (T). The operator
Rº: L' (12)
II TETY
LT(T),
f (R4
(11)
is called the Radon Transform on Tree Systems. This operator is injective.
Tree-Sliced Wasserstein Distance. Tran al. [81] proposed the Tree-Sliced Wasserstein Distance on Systems of Lines (TSW-SL), and later Tran et al. [80] proposed the Distance-batt Tree-Sliced Wasserstein Distance (Db-TSW) which is the generalization of the former. Throughout this paper, we refer to both variants collectively as the Tree-Sliced Wasserstein (TSW) distance for brevity. This notion is distinct from the original TSW distance proposed in [46, 42, 34, 48]. which was primarily developed for applications involving static-support measures, such as classification or topological data analysis. In contrast, the TSW-SI and Db-TSW formulations cast as OT problems over free systems-ard are specifically designed to handle applications with dynamic-suppon measures, as commonly encountered generative modeling tasks. Given μ E P(TRd) and fa. fo are the probability density functions of μr, respectively. The TSW distance between μl, v is defined by
TSW = War. (R dT()),
(12)
TSW is metric on P(Rd). Leveraging the closed-form solution (2) and the Monte Carlo method, TSW in Equation (12) can be efficiently approximated by a closed-form expression.
3 Tree-Sliced Entropy Partial Transport
In this section, we formally introduce the Tree-Sliced Entropy Partial Transport framework and undertake study of its theoretical properties and associated computational complexity.
3.1 Tree-Sliced Entropy Partial Transport
Start with density function f e L (R.). The Radon Transform Ra maps f to density function on a tree system while preserving its total mass, i.e..
I - f f(x) do = | for all TET.
(13)
A proof of this property is provided in Appendix D.6. For u.v E M(Rd), recall that fa and fo denote their respective density functions. Given a tree system T E T and a splitting map C(R" Ty Ak-1), the Radon Transform R", as defined in Equation (10), maps fa and fo to Rt and R9 respectively- measures by μT, PT M(T). We then compute the regularized transport cost d,, (IT, PT) as in Equation (10). The proposed discrepancy is defined as the expectation of this quantity over the space of tree systems T with respect to the sampling distribution
of which are density functions on T. Denote the respective
5

---

Definition 3.1 (Tree-Sliced Entropy Partial Transport), The Tree-Sliced Entropy Partial Transport, denoted as PartialTSW, between μ and v in (RR)) is defined by
(H, = J do (T)
(14)
Remark 3.2. It is worth noting that the value of Db- TSW is determined by range of modeling choices, including the tree system space T, the sampling distribution σ, the splitting map a E C(R TV Ak-1), and the regularization parameters b, X, w(-), and a involved in the definition of da(s)). These dependencies are excluded from the notation for simplicity and readability.
3.2 Properties of PartialTSW
Consider the space I equipped with the Euclidean norm |l- For any vector E R the translation by is the map R² - F defined by x 2 + v. The translation group T(d) consists of all such translations and is isomorphic to the additive group Rdd The orthogonal group ((d) is the group consists of all x d orthogonal matrices. The Euclidean group E(d) comprises all transformations of R that preserve pairwise Euclidean distances. Formally, B(d) is the semidirect product of T(d) and (dd). Each element and acts on IRd via y to the space of tree systems T through the rule
B(d) can be represented as a pair g (Q, v). where Q E O(d) and E R Qy + v. The canonical group action of E(d) on 1 naturally extends
gT {gl. = (Qx; + a, Q0:) E TM.
which preserves the underlying tree structure by construction. A splitting map E C(R x TV As-1) said be B(ddrinvariaat if
algz, gT) a(x T), for all I and T TV-
(15)
In the context of optimal transport theory, where cost function defined on the ground space is lifted distance between measures, it is often desirable- particularly for measures on Rhathat the resulting metric be equivariant under the action of the Euclidean group. Notably, both the Wasserstein distance and the Sliced Wasserstein distance are known to exhibit E(d)- invarrinnei Remarkably, the case of PSrtialTSW, this invariance not only ensures that the discrepancy is E(d)-invariant, but also guarantees that is valid metric on the space of measures M(R.).
Theorem 3.3. PSs is E(d) E(d)-invariant metric M
The proof for Theorem 3.3 is presented in Appendix $D.7. Remark 3.4. As in [80], for the experiments in Section 4, choose the splitting map a such that
(fin BER
[|c (xr + 10:))12 } for all E Rd and T E TT,
alx, T) = softmax
(16)
which is E(d)-invariant.
3.3 Computation of Tree-Sliced Entropy Partial Transport
Similar to UOT and PT, PartialTSW compares μr, v E M(R") while offering a mechanism to softly enforce the fraction of mass to be transported. This is achieved by adjusting the total masses of their respective projections onto a tree system T, denoted ((T) and v(T) (as in Equation (8) and (9)). Given that the Radon Transform R2. is mass- passeprrss.,ing. modifying the masses μ(T) and u(T) directly controls the degree of partiality in the transport between the original measures and v. In practice, (T( is often normalized to unit mass, with (TT) then serving tunable hyperparameter to control this partiality.
A key application of PSrt is computing gradients of PartialTSW(J, v) with respect to samples from μ and v, crucial for generative modeling and gradient flows. Here, samples typically have constant mass (e.g., 1/n where is the number of supports), rendering gradients with respect to total input masses (R) and v(Rd) less meaningful. This simplifies hyperparameter selection: parameters a, b and À in Equation (9) are not necessary. The mass uTT) emerges as the only parameter for controlling the degree of partiality.

---

10a
SOPT
SPOT
Sinkhorn
USOT
SUOT
PAWL
Par
10a
Rundie (s)
10-4
10
10'
10a
Number of Samples
Figure 2: Runtime comparison for PartialTSW and PT/UOT solvers over
Computationally, the intractable integral in Equation (50) is approximated using Monte Carlo method:
(4, = L Σ da(uT,pr.))
(17)
where (Tilf= are tree systems independently sampled from a over T. The theoretical complexity + Lkdn) (with samples, lines per tree, d data dimension), identical to its
O(Ln log log balanced counterpart, TSW [80]. PartialTSW adds negligible computational overhead to TSW, mainly from adjusting the masses v(T). Detailed algorithms and complexity analysis are provided in Appendices SE.1 and $E.2.
The log-linear scaling of PartialTSW with respect the number of samples makes it significantly faster in practice than existing UOT and PT methods. Figure 2 illustrates this performance gap: for 10 samples, alternative approaches, such the rranslaton-invariian Sinkhorn [74] and PAWL [13], are approximately three orders of magnitude slower than PartialTSW. While USOT [7] and SUOT [7] exhibit similar scaling behavior due to their GPU-friendly implementations. Par- tialTSW maintains speed advantage. Further details on computational efficiency are available in Appendix $E.6.
Since PSt is approximated via Monte Carlo (MC) estimation, a crucial aspect the stability and sample complexity of its estimator. Our distance is computed over L sampled trees, and its approximation error is theoretically expected to decrease at standard (LL21/2) rate. We provide detailed empirical analysis in Appendix $E.4 that verifies this convergence.
Experimental Results
This section details the empirical evaluation of PartialTSW against other methods in tasks requiring noise robustness for point cloud alignment, outlier rejection in generative modeling, and effective handling of class imbalance in image image translation.
4.1 Noisy Point Cloud Gradient Flow
In this experiment, we aim to evaluate the robustness of PartialTSW compared to other optimal transport (OT) variants such as SW [10] and Db-TSW [80]. The source dataset X consists of 10,000 points arranged in the shape of dragon. The target dataset Y, also with 10,000 points, follows the shape of bunny and perturbed with 7% noise. Our objective to determine whether PartialTSW can align with the target shape while ignoring noisy outliers. The clean data is taken from [2].
7

---

SS
PartaiTSW 08--33
Figure 3: Visualization of point cloud gradient flows for SW, DS-TSW, and PartialTSW at steps 100, 200, and 300. The leftmost column the source point cloud, and the rightmost column the noise-perturbed target point cloud.
We apply gradient descent to the source points in order to minimize the distance D(X,Y) where D is either PartialTSW. Db-TSW, or SW. The results illustrated in Figure 3. It is clear that OT methods such as SW and Db-TSW are affected by noisy data. In contrast, PartialTSW demonstrates greater robustness, producing smoother interpolations. We refer Appendix $E.7 for further details.
4.2 Robust Generative Model
To further demonstrate its outlier robustness, PartialTSW evaluated generative modeling experiment. First, an Autoencoder pre-trained MNIST digits provides 2D latent representations for digit 0 (the target class) and digit 1 (the outlier class), scaled to approximately reside within [-1, 1 Let X, and X2 denote the true latent distributions for digits 0 and 1, respectively. The generator is subsequently trained using an observed dataset, Xobs. which is a mixture composed of 90% samples drawn from Xg and 10% samples (outliers) drawn from X1. This contaminated imput data is illustrated in Figure 4a. A generator G N (0, I2) [-1, 1]2 is then trained by minimizing D(G (X((),X)))))) Xobs), where Z is batch of noise samples and D is the (Partial) Optimal Transport distance. The objective is for G' to learn to capture the target distribution Xú from the contaminated X effectively ignoring the outliers from X1 Further experimental details are available Appendix $E.8.
Figure and Table summarize method performances. Standard OT methods (e.g., SW [10], Db- TSW [80]: Figure 4b-c) and several UOT/PT approaches (e.g.. SOPT [2], SPOT [9], SUOT [7], USOT [7]: Figure 4d-e, g-h) struggled with the 10% outliers. Despite careful hyperparameter tuning (details in Appendix $E.8.2), these methods often produced mixed Os and or noisy outputs.
Conversely, PartialTSW demonstrates excellent robustness, achieving 0.00% outlier rate (Table 1) by ignoring MNIST 1 outliers and generating only high-quality 0 digits (Figure 4j). This robust performance is complemented by strong sample diversity, stemming from its notably well-distributed latent space. Sinkhorn [74] and PAWL [13] also achieve 0.00% outlier rejection (Figure 4f,i); however, they exhibit less sample diversity, as suggested by their concentrated latent clusters.
Par also provides this robust and diverse generation with high computational efficiency. Its 55s runtime matches Db-TSW and is significantly faster than Sinkhorn (358s) and PAWL. (88s). While SW is the fastest overall, it lacks the necessary outlier robustness.

---

Table I: Quantitative comparison for robust generative modeling MNIST, using target digit 0 and 10% 1 outliers. PartialTSW achieves perfect outlier rejection alongside competitive runtime.
OT
Unisulanced OT Partial OT
Ours
Metric
SW [10)
D6-TSW [80)
SOPT [2)
SPOT
Sinkhorn [74]
SUOT
USOT
PAWL [13]
Partia/TSW
Outliers (%) 1
15.06
16.44
13.28
41.00
17.08
0.00
Runtime (8) 1
37
55
278
278
306
55
(a) Datasct
(b) SW [10)
(c) TSW [80]
(d) SOPT [2)
(c) SPOT [9]
0000010002
(I) Sinkhorn [74]
(g) SUOT [7]
(h) USOT [7]
(i) PAWL [13]
(i) PartialTSW
Figure 4: Qualitative comparison of outlier robustness in generative modeling. Methods learn generate MNIST 0 (blue) from dataset with 10% 1 outliers (orange). Subplots display generated latent distributions (top) and images (bottom). PartialTSW successfully ignores outliers and accurately leamns the true latent distribution of the digit 0, leading to the diverse generated images.
Thus, PS uniquely combines strong outlier rejection, diverse sample generation, and high computational efficiency, making it highly effective for generative modeling in contaminated settings.
4.3 Imbalance Image to Image Translation
Another key attribute of PartialTSW is its ability to handle class imbalance. We demonstrate this capability in an image-to-image translation task, converting "Young" to "Adult" faces using the FFHQ dataset [37]. This dataset presents significant imbalance, with 38K "Young" images and 10.5K "Adult" images. Our experimental setup follows the protocol from recent unbalanced optimal transport studies [22, 28]. Specifically, a pre-trained ALAE autoencoder [64] yields 5 12-dimensio

---

0.52
X. SW
D615W
UOIUM Maght
Table 2: Quantitative Young-to-Adult translation results.
PartialTSSW (Ours) achieves a good balance of Accuracy 1 and perceptual similarity (LPIPS l).
0.32
Method
Accuracy (%6 1
LPIPS 1
SW [10]
89.88 + 0.01
0,4074 + 0,0002
0.44
DS- (80)
89.96 + 0.01
0.4068 +: 0.0002
À 0.05
76.16 + 0.08
0,4979 + 0.0001
0.42
UOT-FM (22)
À = 0.1
82.38 + 0.05
0.4713 = 0.0001
30
0.42
0.2
80.44 + 0.03
0.4920 = 0.0003
50
70.39 + 0.09
0.3706 = 0.0003
--
ULightOT [28)
250
1000
85.27 + 0.06
82.70 + 0.05
0.4086 = 0.0002
0.4198 0.0001
10000
86.06 + 0.04
0.4233 = 0.0001
Accuracy
P(T) L.1
85.71 1: 0.03
0,4047 0.0002
PartialTSW
M(T) 0.9
91.06 + 0.02
0,4183 = 0,0002
Figure 5: Visualizing the Accuracy-LPIPS
M(T) = 0.5
97.05 + 0.03
0.4590 = 0.0001
trade-off in saaagetoo-mmage translation.
M(T) = 0.3
99.11 + 0.03
0,4902 = 0.0003
latent image representations, where the translation is performed. All methods are evaluated on two criteria: (1) translation accuracy (whether images reconstructed from M classify as "Adult") and (2) perceptual similarity (LPIPS [90]) between original and translated reconstructions.
The mapping network M is trained to translate latent samples X from the "Young" domain to align with Y from the "Adult" domain by minimizing distance (XM(X), Y), where D represents PartialTSW, SW [10], or Db-TSW [80]. We also compare against other recent procedures for this task: UOT-FM [22] and ULightOT [28]. These methods, along with PartialTSW, feature parameters to control their degree of regularization when handling data discrepancies: PartialTSW controls the transported mass via its parameter v(T): UOT-FM [22] uses its parameter À to influence the regularization of marginal constraints; and ULightOT [28]'s parameter governs the extent of mass conservation. The adjustment of these parameters creates an Accuracy-LPIPS trade-off (see Figure 5).
Table 2 demonstrates the favorable balance of PartialTSW. Specifically, PartialTSW (with ((T) = 0.9) achieves higher accuracy of 91.06%, versus UOT-FM's (with À 0) 32% and ULightOT's (with v = 10000) 86.06% Concurrently, it achieves a lower LPIPS of 0.4183 (indicating better perceptual similarity), compared UOT-FM's 0.4713 and ULightOT's 0.4233. Furthermore, Par- tillTSW with v(T) = 0.3 attains the highest translation accuracy of 99.11 1111% %, significantly surpassing the accuracies of standard OT methods like SW (89.88%) and Db-TSW (89.96%).
These findings underscore PartialTSW's capability to handle significant class imbalances in image translation, offering solution that effectively balances high target domain alignment with the preservation perceptual similarity. Further experimental details are available in Appendix $E.9.
5 Conclusion
In this paper, we introduce Tree-Sliced Entropy Partial Transport (PartialTSW), a novel distance developed by integrating Entropy-Regularize, Partial Transport for unbalanced measures on tree metric spaces with the Tree-Sliced Wasserstein (TSW) framework on tree systems. We investigate the theoretical properties of the proposed distance and establish that it constitutes valid metric on the space of measures in Euclidean spaces. PartialTSW maintains the computational complexity of the balanced TSW distance, despite being tailored handle unbalanced measures. Crucially, is demonstrably faster than existing Unbalanced and Partial Optimal Transport approaches. Furthermore, comprehensive experiments demonstrate its effectiveness in addressing noise and imbalance in real- world data scenarios. A notable limitation of TSW of existing TSW variants-is is that it defines a distance without yielding explicit transport maps. An important direction for future work therefore constructing optimal partial transport plans within the tree-sliced setting.
10

---

Acknowledgments and Disclosure of Funding
We thank the area chairs and anonymous reviewers for their comments. TL. gratefully acknowledges the support of JSPS KAKENHI Grant number 23K11243, and Mitsui Knowledge Industry Co., Ltd. grant. acknowledges support from the Application Driven Mathematics Program funded and organized by the Vingroup Innovation Fund and VitB..
This research project is supported by the National Research Foundation Singapore under the Al Singapore Programme (AISG Award No:
This research / project is supported by the Ministry of Education, Singapore, under the Academic Research Fund Tier (FY2023) (A-8002040-00-00) A 1-8002039-00-000 This research / project is also supported by the NUS Presidential Young Professorship Award (A 01-0009807001-40 and the NUS Artificial Intelligence InstituteSeed Funding (A (A-8003062-00-00)
References
[1] David Alvarez-Melis, Tommi Jaakkola, and Stefanie Jegelka. Structured optimal transport. In International Conference on Artificial Intelligence and Statistics, pages 1771-1780. PMLR, 2018.
[2] Yikun Bai, Bernhard Schmitzer, Matthew Thorpe, and Soheil Kolouri. Sliced optimal partial transport. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 13681-1 16681- 2023.
[3] Tadas Baltrušaitis, Chaitanya Ahuja, and Louis-Philippe Morency. Multimodal machine learn- ing: survey and taxonomy. IEEE transactions on pattern analysis and machine intelligence, 41(2):423-443, 2018.
[4] Erhan Bayraktar and Gaoyue Guo. Strong equivalence between metrics of Wasserstein type. Electronic Communications in Probability, 26(none):1 - 13, 2021.
[5] Clément Bonet, Paul Berg, Nicolas Courty, François Septier, Lucas Drumetz, and Minh-Ta Tan Pham. Spherical sliced-Waas Wasserstein. arXiv preprint (rXi::2206.08780, 2022.
[6] Clément Bonet, Lactitia Chapel, Lucas Drumetz, and Nicolas Courty. Hyperbolic sliced- Wasserstein via geodesic and horospherical projections. Topological, Algebraic and Geomet- ric Learning Workshops 2023, pages 3340 PMLR, 2023.
[7] Clément Bonet, Kimia Nadjahi, Thibault Sejounne, Kilian Fatras, and Nicolas Courty. Slicing unbalanced optimal transport. Transactions on Machine Learning Research, 2024.
[8] Clément Bonet, Lactitia Chapel, Lucas Drumetz, and Nicolas Courty. Hyperbolic sliced- Wasserstein via geodesic and horospherical projections. In Proceedings of 2nd Annual Workshop on Topology, Algebra, and Geometry in Machine Learning (TAG-ML), pages 334-370. PMLR, 2023.
[9] Nicolas Bonneel and David Cocurjolly. SPOT: sliced partial optimal transport. ACM Transac- tions on Graphics (TOG), 38(4):1-13, 2019.
[10] Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and Radon Wasser- stein barycenters of measures. Journal of Mathematical Imaging and Vision, 51:22 45, 2015.
[11] Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal opti- mal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pages 6511-6522 PMLR, 2022.
[12] Luis A Caffarelli and Robert McCann. Free boundaries in optimal transport and Monge- Ampere obstacle problems. Annals of mathematics, pages 673-730, 2010.
[13] Laetitia Chapel and Romain Tavenard. One for all and all for one: Efficient computation of par- tial Wasserstein distances on the line. In International Conference on Learning Representations, 2025.
11

---

[14] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing systems, 2018.
[15] Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François- -aviier Vialard. Unbalanced optimal transport: Dynamic and kantorovich formulations. Journal of Functional Analysis, 2018.
[16] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
[17] Pinar Demetei, Rebecca Santorella, Manav Chakravarthy, Bjorn Sandstede, and Ritambhara Singh. Scotv2: Single-cell multiomic alignment with disproportionate cell-type representation. Journal of Computational Biology, 29 (0(1)):1213-2288 2022.
[18] Ishan Deshpande, Yuan-Ting Hu, Ruoyu Sun, Ayis Pyrros, Nasir Siddiqui, Sanmi Koyejo, Zhizhen Zhao, David Forsyth, and Alexander G Schwing. Max-sliced Wasserstein distance and its use for GANs. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10648-10656, 2019.
[19] Weitao Du, He Zhang, Yuanqi Du, Qi Meng, Wei Chen, Nanning Zheng, Bin Shao, and Tie- Yan Liu. Se (3) equivariant graph neural networks with complete local frames. In International Conference Machine Learning, pages 5583-5608. PMLR, 2022.
[20] Richard Mansfield Dudley. The speed of mean Glivenko-Cantell convergence. The Annals of Mathematical Statistics, 40(1):40-50 1969.
[21] Steven N Evans and Frederick A Matsen. The phylogenetic Saatorovich-Rubisstee metric for environmental sequence samples. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 74(3):569- 59209995992222 2012.
[22] Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian J Theis. Unbalancedness in neural Monge maps improves unpaired domain translation. In The Twelfth International Conference on Learning Representations, 2024.
[23] Jiaojia Fan, Isabel Haasler, Johan Karlsson, and Yongxin Chen. On the complexity of the optimal transport problem with graph-structured cost. In International Conference on Artificial Intelligence and Statistics, pages 9147-9165. PMLR, 2022.
[24] Kilian Fatras, Thibault Sejourne, Rémi Flamary, and Nicolas Courty. Unbalanced minibatch optimal transport; applications to domain adaptation. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 3186-3197. PMLR, 18-24 Jul 2021.
[25] Kilian Fatras, Younes Zine, Rémi Flamary, Remi Gribonval, and Nicolas Courty. Learning with minibatch Wasserstein: asymptotic and gradient properties. In Silvia Chiappa and Roberto Calandra, editors, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pages 2131-2141, Online, 26-28 Aug 2020. PMLR.
[26] Charles Fefferman and Scott Markel. Recovering feed-forward net from its output. Advances in neural information processing systems, 6, 1993.
[27] Alessio Figalli. The optimal partial transport problem. Archive for rational mechanics and analysis, 195(2):533 55,60, 2010.
[28] Milena Gazdieva, Arip Asadulaev, Evgeny Burnaev, and Alexander Korotin. Light unbalanced optimal transport. In The Thirth-eighth Annual Conference on Neural Information Processing Systems, 2024.
[29] Ziv Goldfeld and Kristjan Greenewald. Sliced mutual information: A scalable measure of statistical dependence. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 17567-17578. Curran Associates, Inc., 2021.
12

---

[30] Elisenda Grigsby, Kathryn Lindsey, and David Rolnick. Hidden symmetries of relu networks. In International Conference on Machine Learning, pages 11734-11760. PMLR, 2023.
[31] Allen Hatcher. Algebraic topology. Cambridge University Press, 2005.
[32] Xinru Hua, Truyen Nguyen, Tam Le, Jose Blanchet, and Viet Anh Nguyen. Dynamic flows on curved space generated by labeled data. In Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence, IJCAI-23, pages 3803-3811, 2023.
[33] Piotr Indyk. Algorithmic applications of low-distortion geometric embeddings. In Proceedings 42nd IEEE Symposium on Foundations of Computer Science (FOCS), pages 10-33, 2001.
[34] Piotr Indyk and Nitin Thaper. Fast image retrieval via embeddings. In International workshop on statistical and computational theories of vision, volume 2, page 5, 2003.
[35] Hoang Anh Just, Feiyang Kang, Tianhao Wang, Yi Zeng, Mycongseob Ko, Ming Jin, and Ruoxi Jía. LAVA: Data valuation without pre-specified learning algorithms. In The Eleventh International Conference on Learning Representations, 2023.
[36] loannis Kalogeropoulos, Giorgos Bouritsas, and Yannis Panagakis. Scale equivariant graph metanetworks. Advances in neural information processing systems, 37:106800-106844) 2024.
[37] Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4401-4410, 2019.
[38] Samuel Kessler, Tam Le, and Vu Nguyen. SAVA: Scalable learning-agnostic data valuation. In The Thirteenth International Conference on Learning Representations, 2025.
[39] Soheil Kolouri, Kimia Nadjahi, Umut Simsekli, Roland Badeau, and Gustavo Rohde. Gener- alized sliced Wasserstein distances. Advances in neural information processing systems, 32, 2019.
[40] Stanislav Kondratyev, Léonard Monsaingeon, and Dmitry Vorotnikov, A fitness-driven cross- difffusion system from population dynamics as a gradient flow. Journal of Differential Equations, 261(5):2844-2088 2016.
[41] Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev, Light Schrödinger bridge. In The Twelfth International Conference Learning Representations, 2024.
[42] Tam Le and Truyen Nguyen. Entropy partial transport with tree metrics: Theory and practice. In International Conference Artificial Intelligence and Statistics, pages 3835 3843-3843, PMLR, 2021.
[43] Tam Le, Truyen Nguyen, and Kenji Fukumizu. Scalable unbalanced Sobolev transport for measures on a graph. In International Conference on Artificial Intelligence and Statistics, pages 8521-8560. PMLR, 2023.
[44] Tam Le, Truyen Nguyen, Hideitsu Hino, and Kenji Fukumizu. An efficient Orlicz-Sobolev ap- proach for transporting unbalanced measures on a graph. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025.
[45] Tam Le, Truyen Nguyen, Dinh Phung, and Viet Anh Nguyen. Sobolev transport: A scalable metric for probability measures with graph metrics. In International Conference on Artificial Intelligence and Statistics, pages 9844-9866 PMLR, 2022.
[46] Tam Le, Makoto Yamada, Kenji Fukumizu, and Marco Cuturi. Tree-sliced variants of Wasser- stein distances. Advances neural information processing systems, 32, 2019.
[47] Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transpon problems and new Hellinggr-KKanorroone distance between positive measures. Inventiones mathemati- cae, 21112):9699911111 2018.
13

---

[48] Ya- Wei Eileen Lin, Ronald R. Coifman, Gal Mishne, and Ronen Talmon. Tree- Fre-Wasserstein distance for high dimensional data with a latent feature hierarchy. The Thirteenth International Conference on Learning Representations, 2025.
[49] Manh Luong, Khai Nguyen, Nhat Ho, Gholamreza Haffari, Dinh Phung, and Lizhen Qu. Revis- iting deep audio-text retrieval through the lens of transportation. In The Twelfth International Conference Leurning Representations, 2024.
[50] Gonzalo Mena and Jonathan Weed. Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem. In Advances in Neural Information Processing Systems, pages 4541-455 15511, 2019.
[51] Boris Muzellec, Julie Josse, Claire Boyer, and Marco Cuturi. Missing data imputation using optimal transport. In International Conference on Machine Learning, pages 7130-7140. PMLR, 2020.
[52] Kimia Nadjahi, Alain Durmus, Lénaïc Chizat, Soheil Kolouri, Shahin Shahrampour, and Umut Simsekli. Statistical and topological properties of sliced probability divergences. Advances in Neural Information Processing Systems, 33:20802-20812, 2020.
[53] Khai Nguyen, Shujian Zhang, Tam Le, and Nhat Ho. Sliced Wasserstein with random-path projecting directions. In Forty-first International Conference on Machine Learning, 2024.
[54] Tin D Nguyen, Brian L Trippe, and Tamara Broderick. Many processors, little time: MCMC for partitions via optimal transport couplings. In International Conference on Artificial Intelligence and Statistics, pages 3483-3514. PMLR, 2022.
[55] Trung Nguyen, Quang- Hieu Pham, Tam Le. Tung Pham, Nhat Ho, and Binh-Son Hua. Point-set distances for learning representations of 3d point clouds. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 10478-10487, 2021.
[56] Vu Nguyen, Tam Le, Makoto Yamada, and Michael A Osborne. Optimal transport kernels for sequential and parallel neural architecture search. In International Conference on Machine Learning, pages 8084-8095. PMLR, 2021.
[57] Sloan Nietert, Ziv Goldfeld, and Rachel Cummings. Outlier-robust optimal transport: Duality, structure, and statistical analysis. In International Conference on Artificial Intelligence and Statistics, pages 11691-11719, PMLR, 2022.
[58] Jonathan Niles-Weed and Philippe Rigollet. Estimation of Wasserstein distances in the spiked transport model. Bernoulli, (644))66662688. 2022.
[59] Jungin Park, Jiyoung Lee, and Kwanghoon Sohn. Bridging vision and language spaces with assignment prediction. In The Twelfth International Conference on Learning Representations, 2024.
[60] François-Pierre Paty and Marco Cuturi. Subspace robust Wasserstein distances. In Proceedings of the 36th International Conference on Machine Learning, pages 5072-5081, 2019.
[61] Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trendy@ in Machine Learning, 11(5-6):355-677, 2019.
[62] Khiem Pham, Khang Le, Nhat Ho, Tung Pham, and Hung Bui. On unbalanced optimal transport: An analysis of Sinkhorn algorithm. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 7673-7682. PMLR, 13-18 Jul 2020.
[63] Thong Pham, Shohei Shimizu, Hideitsu Hino, and Tam Le. Scalable counterfactual distribution estimation in multivariate causal models. In Conference on Causal Learning and Reasoning (CLeaR), 2024.
[64] Stanislav Pidhorskyi, Donald A Adjeroh, and Gianfranco Doretto. Adversarial latent autoen- coders. In Proceedings of the IEEF/CVE conference on computer vision and pattern recognition, pages 14104-141 1111,, 2020.
14

---

[65] Michael Quellmalz, Robert Beinert, and Gabriele Steidl. Sliced optimal transport on the sphere. Inverse Problems, 39(10):105005, 2023.
[66] Julien Rabin, Gabriel Peyré, Julie Delon, and Mare Bemot. Wasserstein barycenter and its application to texture mixing. In International Conference on Scale Space and Variational Methods in Computer Vision, pages 435-446, 2011.
[67] Julien Rabin, Gabriel Peyré, Julie Delon, and Mare Bemot. Wasserstein barycenter and its application to texture mixing. In Scale Space and Variational Methods Computer Vision: Third International Conference, SSVM 201 I, Ein-Gedi, Israel, May 29-June 2, 2011, Revised Selected Pupers 3. pages 435 Springer, 2012.
[68] Gathika Ratnayaka, James Nichols, and Qing Wang. Learning partial graph matching via optimal partial transport. In The Thirteenth International Conference on Learning Representations, 2025.
[69] Grant Rotskoff, Samy Jelassi, Joan Bruna, and Eric Vanden-Eijnden. Global convergence of neuron birth-death dynamics. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Pro- ceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 5508-5517, 2019.
[70] Mahdi Saleh, Shun-Cheng Wu, Luca Cosmo, Nassir Navab, Benjamin Busam, and Federico Tombari. Bending graphs: Hierarchical shape matching using gated optimal transport. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 11757-11767. 2022.
[71] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference machine learning, pages 9323 93329. PMLR, 2021.
[72] Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal transsoot analysis of single-cell gene expression identifies developmental trajectories reprogramming. Cell, 176(4):928-94) 943, 2019.
[73] Thibault Séjouréé, Gabriel Peyré, and François Xavier Vialard. Unbalanced optimal transport, from theory to numerics. arXiv preprint arXiv:2211.08775: 2022.
[74] Thibault Sejourne, Francois- -aaccis-Xavvie Vialard, and Gabriel Peyré. Faster unbalanced optimal transport: Translation invariant Sinkhorn and 1-D Frank-Wolfe. Wolfe. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151, pages 4995-5021. PMLR, 2022.
[75] Charles Semple and A Mike. Steel, phylogenetics, vol. 24. Oxford Lecture Series in Mathematics and its Applications. Oxford University Press, Oxford, 2003.
[76] Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional Wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphics (TOG), 34(4):66, 2015.
[77] Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto- encoders. In International Conference on Learning Representations, 2018.
[78] Hoang Tran, Thieu Vo, Tho Huu, Tan Nguyen, et al. Monomial matrix group equivariant neural functional networks. Advances in Neural Information Processing Systems, 77:48628-48666 2025.
[79] Hoang V. Tran, Thanh Chu, Minh-Khoi Nguyen-Nhat, Huyen Trang Pham, Tam Le, and Tan Minh Nguyen. Spherical tree-sliced Wasserstein distance. In The Thirteenth International Conference on Learning Representations, 2025.
[80] Hoang V. Tran, Minh- Khoi Nguyen-Nhat, Huyen Trang Pham, Thanh Chu, Tam Le, and Tan Minh Nguyen. Distance- based tree-sliced Wasserstein distance. In The Thirteenth Interna- tional Conference on Learning Representations, 2025.
15

---

[81] Hoang V. Tran, Huyen Trang Pham, Tho Tran Huu, Minh-Khoi Nguyen-Nhat Thanh Chu, Tam Le, and Tan Minh Nguyen. Tree-sliced Wasserstein distance: A geometric perspective. In Forty-second International Conference on Machine Learning. 2025.
[82] Hoang- Viet Tran, Thieu N Vo, Tho Tran Huu, and Tan Minh Nguyen. A clifford algebraic ap- proach to e (u)equvivriiant high-order graph neural networks. arXiv preprint arXiv:2410.04622 2024.
[83] Huy Tran, Yikun Bai, Abihith Kothapalli, Ashkan Shahbazi, Xinran Liu, Rocio P Diaz Mar- tin, and Soheil Kolouri. Stereographic spherical sliced Wasserstein distances. In Forty-first International Conference on Machine Learning, 2024.
[84] Thanh Tran, Hoang V. Tran, Thanh Chu, Huyen Trang Pham, Laurent Ghaoui, Tam Le, and Tan Minh Nguyen. Tree-sliced Wasserstein distance with nonlinear projection. In Forty-second International Conference on Machine Learning. 2025.
[85] Viet-Hoang Tran, Van Hoan Trinh, Khanh Vinh Bui, and Tan M Nguyen. On linear mode connectivity of mixture-of-expert. architectures. arXiv preprint arXiv:2509.11348 2025.
[86] Viet-Hoang Tran, Thicu N Vo, An Nguyen The, Tho Tran Huu, Minh-Kha Khoi Nguyen-Nul -hhat, Thanh Tran, Duy-Tung Pham, and Tan Minh Nguyen. Equivariant neural functional networks for transformers. arXiv preprint 24. 2024.
[87] C. Villani. Optimal Transport: Old and New, volume 338. Springer Science & Business Media, 2008.
[88] Thieu N Vo, Viet-Hoang Tran, Tho Tran Huu, An Nguyen The, Thanh Tran, Minh-Khi Nguyen- Nhat, Duy-Tung Pham, and Tan Minh Nguyen. Equivariant polynomial functional networks. arXiv preprini arXiv:24 241.04221 2024.
[89] Anh-Khoa Nguyen Vu, Toan Do. Vinh-Tiep Nguyen, Tam Le, Minh-Triet Tran, and Tam Nguyen. Few-shot object detection via synthetic features with optimal transport. Computer Vision and Image Understanding, page 104350, 2025.
[90] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreason- able effectiveness of deep features as perceptual metric. In CVPR, 2018.
[91] Bo Zhao, Robert M Gower, Robin Walters, and Rose Yu. Improving convergence and general- ization using parameter symmetries. arXiv preprint arXiv:2305.11 13404, 2023.
16

---

Table of Notation
R
d-dimensional Euclidean space
Euclidean norm standard dot product
gd
(d 1)-dimensional hypersphere
0
unit vector
disjoint union
L'(X)
space of Lebesgue integrable functions on X
P(X)
space of probability measures on X
M(X)
space of measures on X
p.
measures
8(.)
1-dimensional Dirac delta function
u((
uniform distribution on g
d
pushforward (measure)
C(X,))
space of continuous maps from X to Y metric metric space
.(:.))
dr(-.)
tree metric
T(d)
translation group order d orthogonal group of order Euclidean group of order element of group D-Wasserstein distance
((d)
E(d)
g
W
SW, g
Sliced p- Wasserstein distance
(rooted) subtree
edge graph
We
weight of edge in graph
1
line, index of line
system lines, tree system ground set of system of lines, tree system space of systems of lines in Rd tree structure in system of lines
C.
L
T
number of tree systems
k
number lines system of lines or tree system
Ra
Radon Transform Systems of Lines
AR-1
(A: 1)-dimensional standard simplex splitting map tuning parameter in splitting maps space of tree systems
g
distribution space of tree systems
17

---

Appendix of "Tree-Sliced Entropy Partial Transport"
Table of Contents
Background Optimal Transport on Metric Spaces with Tree Metrics
19
Background on Entropy Partial Transport on Metric Spaces with Tree Metrics
19
C Background on Tree-Sliced Wasserstein Distance Euclidean Spaces
22
C.I Tree System
22
C.2 A Variant of Radon Transform for Systems of Lines
23
C,3 Tree-Sliced Wasserstein Distance for Probability Measures on Euclidean Spaces
24
Theoretical Proofs
24
D.1 Proof for Proposition B.I
24
D.2 Proof for Theorem B.2
25
D.3 Proof for Proposition B.3
30
D.4 Proof for Proposition B.4
31
D.5 Proof for Proposition B.5
31
D.6 Proof for Equation (13)
................................
32
D.7 Proof for Theorem 3.3
E
Experimental Details
34
E.1
Algorithm for Partial Tree-Sliced Wasserstein ................................
E,2
Computational and Memory Complexity Analysis
35
E.3
Empirical Runtime and Memory Performance of PartialTSW
35
E,4
Sample Complexity and Estimator Stability
36
E,5
Discussion on Hyperparameters of Evaluated Methods
................................
37
E.6
Comparing Computational Efficiency
37
E.7
Noisy Point Cloud Gradient Flow
38
................................
E.8
Robust Generative Model
................................
39
E.8.1
Implementation detail
39
E.8.2
Ablation result for baselines
40
E.9
Imbalance Image to Image Translation
49
E.9.1
Implementation detail
49
E.9.2
Additional Experimental Results
50
Boarder Impacts
18

---

A Background on Optimal Transport on Metric Spaces with Tree Metrics
7 (V, E3) be a tree rooted at node r, where each edge e E is assigned a nonnegative length We. Here, V denotes the set of nodes and E' the set of edges, For notational convenience, we use T also refer to the union all nodes and the continuous points along the edges. We now recall the formal definition of a tree metric:
Definition A.1 (Tree metric [75, Section 7, p. 145-1822. A metric d Q x 52 [0, oo) is said to be tree metric set Ω if there exists tree T such that Ω C T and, for all x, y E 2 the distance dax, y) equals the length of the unique path T connecting x and y.
Suppose V is subset of a vector space, and let d' denote the tree metric defined on T. We denote by [z. y] the unique shortest path in between any two points x: and y. Let w be the unique Borel (Iength) measure on T satisfying w([z, v)) d (z, v) for all x, y E Т. For any E T, we define the subtree rooted at æ by
A(z) = (y E I
(18)
Let P(T) denote the set of all probability measures on T, i.e., Borel measures with total mass equal to one. The following result provides closed-form expression for the 1-Wasserstein distance on the tree metric space T.
Theorem A.2 (Optimal Transport on Tree Metric Spaces [46, Section 3, Proposition ID). For any U. P(T), the 1- 1-Wasserstein distance with respect to the tree metric dr is given by
W1.d7 (p. v) J Im((()))
v(A(x)) -
w(dz).
(19)
B Background on Entropy Partial Transport on Metric Spaces with Tree Metrics
In this section, we revisit the Entropy Partial Transport (EPT) formulation introduced in [42] for completeness. All theoretical proofs are outlined in Appendix D.
We denote by M(T) the collection of all nonnegative Borel measures on with finite total mass. Let C(T) denote the space of continuous functions defined on T, and let (T) denote the space of Borel measurable functions on T that are essentially bounded with respect to the measure w. The space (T) forms Banach space when equipped with the norm
[ == inf (a R [f(x)) S for w-almost every x E TT.
(20)
Let MTT T) denote the space of all nonnegative Borel measures on T T with finite total mass. Given ,4, / E M(T), define the set of admissible partial couplings as
11 (1, = {r E M(T T) m S p, 72 1)
(21)
where 71 and 7/2 represent the marginals y on the first and second coordinates, respectively.
For any y E IIg (H, v),, let fi and fz be the Radon Nikodym derivatives of 71 with respect to μ and 72 with respect to respectively. That is, 71 fip and 7/2 f22, with the constraints 0 < fi < jı-a.. and ( < f2 < v-a.e.
Let w : T - [0, co) be a b-Lipschitz continuous and nonnegative weight function, defined by
w(x) a₁ d' (z, vo) + ao.
(22)
where xo E T,, aj E [0, b), and as E [0, oo). Here, .) denotes the tree metric over T. We use the entropy function F [0, oo) (0, o0) given by
F(s) = |s 1/-
Letting m. == min(u(T), ( (T)). and fixing m E [0, m].. the EPT problem is formulated as
Film ( F2(12 v) + { dr u) 7dz dy)
Wm(uvv): =
inf rElls (p..) >T =mm=
(23)
19

---

where the regularization terms are defined as the weighted relative entropies
Fi (A) = | w(x) ((lfi(z)) F2(12 = / w(x) F(fe(=)) v(dx).
(24)
To handle the mass constraint 7(T x T) = m, we introduce a Lagrange multiplier À E R and instead consider the relaxed objective
Fi(ma ()) + F2(22 2) + (
(dr(a,)) A) r(zz, dy)
ET (T == inf yell (H.0r)
(25)
TXT
We now expand the entropic terms and define
Cx(7): = f w(x) p(dz)) + / w(x) w(x) v(dx)
/ (o(x) w(x) 22 (2(xx)
+ b TXT
(dT(s,)) - r(de, dy),
(26)
so that Equation (25) equivalent to
ETx(μ,)) (u,))
inf Cx(1).
(27)
As established in [42, Theorem 3.1, part i)], the solutions Equation (23) and Equation (27) are related via the identity
Wm (n, v)) ETX(U,)) v) + Abm.
(28)
Inspired by the construction proposed by Caffarelli and McCann [12], we recast the entropy- regularized partial transport problem in Equation (27) classical optimal transport (OT) problem between balanced measures. To achieve this, we augment the original domain T by introducing an auxiliary point & $ T, and define the extended space T := TU
We then lift the unbalanced measures 14, E M(T) to balanced counterparts supported on t: = + ((T) 82, ù: v ((T) 82,
(29)
where dg denotes the Dirac measure at point 8. Next, we define cost function è: T x + - R that extends the original transport cost:
b (d - X] w(x)
if if y E T if T and y if E T and x = 3, if = s.
a(x, v)
(30)
w(y)
Using this extended cost, we formulate the balanced OT problem over fa and D: f
KT ==
inf (EI(AA,) +xt Ö(x, v) (dxx, dy).
(31)
where the set of admissible transport plans r(A, is given by:
r'a, 0) = {i E MMT x +) 70 x +) A(U), 1( x U) D(U), v Borel set U C c
(32)
The connection between the entropy-regularizee partial transport formulation ET, in Equation (27) and the balanced optimal transport problem KT in Equation (31) is established by the following result.
Proposition B.1 (Equivalence of ETx and KT). Let p.v. M(T). Then the two formulations coincide:
ET v) KT 8).
(33)
Furthermore, the optimal plans y for the partial transport problem and f for the balanced transport problem are related by:
== (1 fi) @ 8, + 8, 0 (1 f2)v 7T x T)) where Ji and f2 are the Ridon-Nikodyy derivatives of the marginals of 7 with respect to p and v, respectively.
(34)
20

---

The detailed proof is provided in Appendix $D.1. Note that KT corresponds to a classical optimal transport problem defined between two balanced measures over the extended space T and governed by the cost function è. This allows us to invoke standard OT duality theory, such as [12, Corollary 2.6]. to obtain variational dual formulation for ET,, described below.
Theorem B.2 (Dual Representation of ET x. The dual problem associated with the entropy- regularized partial transport functiooal ETx v) is given by:
x v) =su { J (dp dv): EIEL,
by (M(T)+u(T)
(35)
2
where the admissible funntionn class L is defined as
L = ( : :-
bd S sf 2
bd
If(x) f(y)l S bdd (z, v) for all x, x v.
w 2
The proof of Theorem B.2 is deferred to Appendix $D.2. To obtain a tractable approximation of the dual problem, introduce a regularization based restricted class of test functions. Let denote the root of the tree T, and let w be the associated length measure on T. For a fixed parameter wr(r)] define the function class The to consist of all functions f - R of the form:
E [0, 2
/ g(y) w(y),
f(x) 8 +
(36)
where is a constant satisfying
E w() - bA
by + a, w(r) +
2 a
(37)
2
and g I is a bounded function with I|gll (T) S b.
The a-regularized entropy partial transport is then defined as:
(Tx(L)) sup sella { f
(dp (elu-dv/) dr) bd lu(T) (T)]. +(T))
(38)
This regularized formulation admits an explicit closed-form expression:
Proposition B.3 (Closed-Form Solution for ET,, For u, E M(T), we have:
ET, J Iu(A(=)) v(A(x)) w(dx)
6\ (m((T) [TT]+ + w((r)
6Λ
a Iu(T)- u(T).
(39)
2
The proof is provided in Appendix $D.3. We now compare the original entropy transport value ETx with its regularized approximation ETx:
Proposition B.4 (Comparison Bounds between ET and ET,). The following inequalities hold:
ETx(H, v) < (U,
(40)
and if the condition
[ALT - Alo s 2w(r), where LT == max ([r,
(41)
2E7
satisfied, then for all a such that 26l Sas
ETx v) < ETx(UØ)
(42)
AA
+ w(r).
The proof appears in Appendix SD.4. For S a < cost is defined as:
h1
+ w(r), recall that the regularized transport
(1, = ET (H, v) + (u(TT + +T])..
(43)
2
This cost function defines a genuine metric, as shown below: Proposition B.5 (Metric Structure of d,). (MM(),, da) is complete metric space.
The proof is presented in Appendix $D.5.
21

---

Figure 6: An illustration of the tree system construction is presented the two-dimensional space R², though the method readily generalizes to higher dimensions. The process starts with collection of infinite lines arranged without any inherent structure. are identified (some of which may not be visible in the figure due to the lines' unbounded nature). A subset intersections selected and marked in red to indicate those to be discarded. The remaining intersections are retained to enforce tree structure on the system- gestem-ensuring that any two points lying on the lines are connected by a unique path that passes only through the preserved intersections. These remaining intersections act the essential nodes that define the tree topology. Once the red (discarded) intersections are removed, the resulting configuration forms the desired tree system.
pairwise intersections among these lines
C Background on Tree-Sliced Wasserstein Distance on Euclidean Spaces
This section reviews foundational concepts underlying the Tree-Sliced Wasserstein distance defined Tree Systems. To ensure completeness, revisit key definitions and theoretical formulations; detailed proofs and additional exposition are available in [81, 80].
C.1
Tree System
Line. A line in the Euclidean space Rd is specified by tuple (x, 0) E RR Sd-1 and is expressed parametrically as z + t 0 for t E R. Throughout, we use = (æ), 0:) E Rd x Sd-1 to denote line, where 2; denotes the source point and 6; the direction vector.
System of lines. Given an integer k > 1, system k lines in gd refers to a collection of k such lines. The notation (Rd gd-1)* is abbreviated as Lg. representing the space of systems of lines specific system of lines.
in Rdd. An element in this space, commonly denoted by L, corresponds
Tree System. A system C is said to be connected if the union of all points lying on the constituent lines forms a connected subset of gd By selectively removing certain intersection points between the lines, one can enforce a tree structure on L --eeldingg tree system-in which any two points are connected by a unique path. An illustration of this construction is provided in Figure 6.
Remark C.1. The term tree system is used because there is unique path between any two points, analogous the definition of trees in graph theory.
Beginning with the remaining intersections, we employ the concepts disjoint union and quotient topology [31] to construct tree system by coherently gluing together multiple copies of E. This
22

---

topological framework induces natural metric, under which the resulting space satisfies the properties of metric space.
Sampling Procedure for Chain-Structured Tree Systems. The space tree systems is inherently rich and diverse, primarily due to the wide range of possible underlying tree topologies. [81] presents a general framework that accommodates arbitrary tree structures, while placing particular emphasis on subclass of chain-like trees. The following describes the sampling procedure for generating tree systems with this chain-based architecture:
Step 1. Draw initial point 41 Ni and direction 01 ~ V1, where μl is probability measure on Id and 1/1 is a measure on the unit sphere gd-1
Step i. For each subsequent node, sample t; μ4 and 0g 14, then compute 24-1 + tq 0z-1- Here, μ4 is a distribution over and V; over $d-1
All distributions μ4; and i are assumed to be mutually independent. Specifically, we consider the following choices: The initial position distribution μl is supported bounded subset of IR such as the uniform distribution over the cube [-1, i.e. 4(1-1, 1]9;; For i > 1, each μui is defined on bounded interval of for example, 41[-1, 1]): Finally, each direction O; is drawn from distribution over the unit sphere, e.g., the uniform distribution ((( An example of such a tree system is illustrated in Figure
Remark C.2. This generative process induces probability measure on the space T of all chain- structured tree systems produced via this construction.
C.2 A Variant of Radon Transform for Systems of Lines
Let L (IRd) denote the space of Lebesgue integrable functions on IR equipped with the standard L' norm || |/1- Consider configuration of i lines L E L. A real-valued function f defined on the domain L consists of all points of L, is said to be integrable over the line system if the following condition holds:
llflle = Σ E/ (f dtz < 00.
(44)
IEC
The set of such functions is denoted by (C), representing the space of Lebesgue integrable functions over the line system Recall the standard (k 1)-sipplex:
A (a: (a))Ec E R a 2 0, Σ 5a=14
(45)
LEC
Define the space C(IR LL Ax-1) 1) as the set of continuous maps from Id L to Ag- referred as splitting maps. Given line system L E Lg and a splitting map a E C(I x Lİ AR-1),
define linear operator that projects function f L (-(R)) the line system as follows:
Ref:C. E
(46)
(x, - J f(y) (()) a)y,C) C) 0tz (y xn, 0.)) dy.
(47)
where d denotes the Dirac delta function in one dimension, and (x1, 0t) encodes the location and direction of line I. It can be shown that R belongs to L (C) for any f E L L'(R",, and furthermore satisfies the inequality
[R S 11/1I-
Hence, the operator R2: (IR)) (2) well-defined. These properties are proven in [80]. Extending this to all line systems, we define the Radon transform on Systems of Lines as follows. For a fixed splitting map a E C(Rd Lgg. A+-1). define:
Rº: I. (TR)
II (2)
(48)
CEL
(Ref)cetg)
(49)
If the splitting map is invariant under the Euclidean group E(d) the group of all isometries of then the operator R2 is injective.
23

---

C.3 Tree-Sliced Wasserstein Distance for Probability Measures on Euclidean Spaces
Let p.v. P(RR) be probability measures. For a structtuu system of lines E and an F.d) iivvarintt splitting map a E C(R2 x L.g AR-1), the transform Rg pushes forward j and to corresponding measures 2 and R on L. Since each L E 'T is equipped with a tree metric de. the I-Wasserstein distance W &g.1 between the transformed measures can be computed. This leads to the following definition of the Distance-based Free-Sliced Wasserstein (Db-TSW) [80] distance:
Db-TSW((, = Wae.1 (REH, REv) do(
(50)
where a is a probability measure over the space of tree systems T. It is important to note that the value of Db-TSW depends on the choice of the tree system space TT, the sampling distribution and the specific E(d))invariant splitting map a, although these dependencies are omitted from the notation for brevity. The resulting Db-TSW defines an ((d)-invariant metric on P(R))).
Remark C.3. As established in [80], if the tree systems consist solely of single line, the Db-TSW distance reduces exactly to the classical Sliced Wasserstein (SW) distance on Rdd
Constructing SIvaaranv Splitting Maps. The Euclidean group E(d) consists of all transforma- tions of @dd that preserve pairwise Euclidean distances. As such, this invariance extends only to distances between points but also to the shortest distance from a point to a line. Given a point x E R and system of lines L L define the distance from
line / by:
d(x, L) inf ||= v112-
(51)
yEl
This quantity is preserved under the action of E(d), meaning that any function constructed solely from the collection ()rx, will inherit B(d)-invariance.
Based on this observation, invariant splitting maps is constructed by applying a post-processing function 3: I 1 to the vector of distances. The resulting splitting map,
alz, C) (fd(z, C))ee)
(52)
guaranteed to be E(d)-invariant for any choice of continuous B. Empirically, effective performance in applications is achieved when 3 is taken to be the softmax function with a tunable scaling parameter & > This yields the practical definition:
a(x, L), softmax ((g d(x,, L)i
(53)
which distributes weights across lines in L according to their proximity to z, while respecting the geometric symmetries of the Euclidean space.
Remark C.4. The concepts of equivariance and invariance are widely employed machine learning to ensure model robustness under transformations that preserve the semantic structural properties of the input. Such principles are foundational in the design of architectures that respect inherent sym- metries within data. Applications of cquivariant models span various domains, including equivariant graph neural networks [71, 19, 82], equivariant metanetworks [86, 78, 88, 36], parameter symmetry [30, 26, 85), and optimization [91], among others.
D Theoretical Proofs
To ensure completeness, we provide full derivations of the result, closely following the methodology of [42].
D.1 Proof for Proposition B.1
To ensure completeness, we provide full derivations of the result, closely following the methodology of [42].
Proof. We begin by proving the inequality KT(A,, 8) S (T, Let y E Tg, v) be any admissible partial transport plan, and define Ý according to the expression in Equation (34). By
24

---

construction, 3 T(A, D.. Then, evaluating the cost of / under the extended transport objective yields:
KT Iaxy ((2, (ddr), dy)
b $ (dr( v) >] r(dx, dy)
1 w(x) [1 f w((x) w()) [2(x)) u(dx). (54)
Taking the infimum over all y E (u, on the right-hand side implies: KT 2) & (1)
(55)
We now establish the reverse inequality, i.e., (T 2) 2 (T, P). Let E T(a, 2) be any feasible coupling in the balanced OT problem, and let be its restriction to T T. Then, by [42, Lemma 3.2], we have y E II (14, v) and the decomposition in Equation (34) holds. We now compute the total cost of under ē:
J ((z,y) (dx,, dy) 6 (dT(s,,) ->] 7(dx, dy)
TXT
+ / w(xx [1 (x)) + / w(x) [1 fa(x))v(dz)
2 ET)
(56)
Taking the infimum over all admissible & E ((ū, yields the desired inequality:
KT D) 2 (T)
(57)
Combining both bounds, we conclude the equivalence:
KT D) (Tx
(58)
The correspondence between optimal couplings y and & follows directly from the construction and identities established above.
D.2 Proof for Theorem B.2
To ensure completeness, provide full derivations of the result, closely following the methodology of [42].
Proof. We begin by establishing the intermediate result:
Ur u(xx) u(dx) + +ʃ u(o))
(T, (u,))E
sup
(59)
where the admissible set K. is defined
: (u, v) (p) (v)
u(x) w(x), Vr ET.
bb+ + inf [bd - w(x)) s v(y) s sw(),, VEE =TT
<(x) + w(v) < b(d7(x.)) A), Vo.yET
This identity follows from the dual representation of KT D) via Proposition B.1 and [12, Corollary 2.6], which yields:
Ut ((x) a(dx) + f 6()) p(dy)
(2,))
sup Ger (M), (er (0) ((s)+ I.
(60)
25

---

We aim to show that this supremum I coincides with
U u(x) u(dx) / v(z) u(dr)) (u,v)e
: sup
(61)
To show / 2 J. let (u, v) E K. Extend these functions to 7 by setting:
{ a(x) if ET, if = 8,
((x) ((x) = ( v()) if T,
ú(x)
i 8.
Since (u, v) K, follows directly from the definition of é that ū(x) + T(y) - c(2,)) y) for all x, F. Consequently:
/ (()) + f ((o)
12 J f u(z) + o(x))(de))
(62)
which implies I 2 J.
To prove the reverse inequality I < J, let (û, 6) be a maximizer for Without loss of generality, we can normalize u(8) 0 by observing that replacing (ū, 6) with (ū =((), + ū({)) preserves admissibility and the objective value. Moreover, define:
vyy) inf (d(1, v) ((x)) Vy f.. xe7
(63)
Then ((v) S v(y), and (û, v) remains admissible and achieves the same supremum, so we may further assume i(y) infa [ y) i()) and (()) = 0. In particular,
0()) inf [0(2 8) a(x)] x
(64)
To proceed, we define w(() 0 and consider two cases based on the structure of ú and é.
Case I. Suppose that
inf (wxx) ú(=)] 2 0.
(65)
=E7
In this case, observe that ((5) 0 by assumption. Since
0(8, i(6) = 0 and
inf [ā(1, 8) ú(<:)] inf [1r(x) -(x)) O,
zET
#ET
we conclude that
6(6) inl [((x, ū(xx)] 0.
(66)
=et
Next, for all y F. bound f(y) from above:
((y) inf [é(x, v) axx)] < ((5, a(() w(v),
(67)
xET
where we have used (() = 0. To lower-bound ((y) for y T, note that u(x) < w(x) for all a E T, and w=() 0. Therefore,
((v) inf [0(x, 9) ú(z)]
> inf [ö(x, y) w(x)]
inf [b(d7 (x,, v) A) -w)]
TET
-- + inf [b (2. v) w(x)
(68)
EET
26

---

Combining both bounds, we find that o(y) satisfies all constraints in the definition of K,, and û S w holds by assumption. Hence, (ū, i) E K. We now compute the dual objective:
I =
ú(x) (dxx) + ((z)
a(x) + / (d)) +08)
f ú(x) p(dx) + f ((z))(dr)
(69)
S J.
(70)
Thus, under this case, the supremum / is bounded above by J, completing the proof for Case Case 2. Suppose now that
inf =
(w(xx) -((x)] O.
(71)
As in Case 1, we deduce that
6(8) inf (w(x) ú(xx]] O,
(72)
*ET
and the dual objective becomes
T= ff (()) + ((x))((x)) + ((T) inf [w - új.
(73)
T
Define truncated version of by setting:
ū(x) == min w(x)).
(74)
This ensures that ū(x)= S w(x) and, since (($) 0, also have i($) 0. Furthermore, for all I, T.
((x) (0) 6(x, v),
(75)
due to the pointwise minimum structure of û and the feasibility of (). Since inLzeT [ur(x) - ú(x)] 0, there exists xo E T such that û(xa) > w(xo). Thus, at zo. we have ū(x)) = w((nn) and therefore
inf [ il s. 0.
(76)
On the other hand, since ü(x)= S w(x) everywhere, it follows that [w 2i] 0.
inf
(77)
+
We now rewrite the first two terms in Equation (73) as:
J ((z) p(dx) + U(T) inf(w w úi = / ū(x)) m(dx) + / (=:0())))(()))
[x(x)) w(x)) u(dx)
T
+ m(T) ins[w ú]
J ū(x) u(dx).
(78)
Substituting this into Equation (73), we obtain the upper bound:
I 1 f ū(x) ū(x) m(dx) + f 6(x) ((x) v(dz).
(79)
We now define a new function 0: T T RR by
ö(v) == inf [0(x, v) ū(x)).
(80)
#
27

---

By construction, i(()) 2 ö(y) and for all E T, jyy) S c(3, y) ú(s) w(y).
(81)
Furthermore, using ö(x) S w(xx and the form of è, we obtain a lower bound:
T(y) inf [0(x, y) ū(x)]
TtT
2 inf [b(d7 (2, (2. v) X) w[w(
TET
= -b + inf [b v) - w(x)
(82)
TET
Combining these, we find that (ū, 8) K. Hence,
I I J ((x) m(dx)) + 1 ((x))(dd) s J.
(83)
This completes the analysis for Case 2 and thus confirms the desired equality:
(futa u(x) u(dx) + / v(x) o(x) u(de) (Tx(M,)) = sup (u,e)EK
(84)
where
1C = (u.u): uu
-b + inf
[b (d. v) w(x)] S v(v) S w(y),
aET
u(x) + v(v) S ((d7 v) - A
(85)
We are ready to complete the proof of the theorem. Since the weight function w is 6-Lipschitz, it satisfies the following inequality for all x T:
- ((a)) < inf (bd7(x, -w))
(86)
gET
Let (u,v) v) IC be arbitrary. Define the following sequence of dual potentials via infimal convolutions:
0" (w) = inf (1, v) >1 v(v)} = -b + inf (bdT(r, v) -vy)] 2 u(x),
(87)
WET
#ET
"n* (y) inf inf (b(dd >] - -- (x)) - inf [bdT(x (2, v) vz(x) 2 vly). zET
(88)
Now, observe that the lower and upper bounds for imply that
-b inf [b y) w(x)] S v(y) S w(y). ZET
Using this together with Equation (86), we can derive pointwise bounds on o" for any x T:
("(x) < bb v(x)= < inf [b (z, v) w(y) < w(x),
(89)
gET
(x) 2 Dλ + inf db (z. y) w(y) 2 -b - w(x).
(90)
NET
We now show that o* is b-1 L-Lisschitz. Let æ,, 22 E T and fix an arbitrary & > O. By the definition of infimum, there exists yn T such that
bd bd 81) v(yn) < 0- (x))) + bA +
Then,
2" (x2)) - (x1) < bdr v1) - v(y1) - (bd -
- v(y1)l +
b [dr v1) d v1)l s bdr +e.
(91)
Since this holds for all E > 0, we conclude that
v² (x(2)) o (x1) < bx
(92)
28

---

By symmetry, the reverse inequality also holds, so
(1)())) -
(x2)) (
bdT(
(93)
which confirms that e" is b-Lipschitz.
Thus, v" belongs the following class of functions:
: ( E C(T): -bÀ w((x)= s f(x)) S w(x), [f(x) fw)) s bdt(z.
(94)
This concludes the key regularity properties needed for the dual formulation.
We now establish the identity V** = -b v² To begin, note from the definition that:
(1) inf (b( v) >) (x) S -b v² vy()..
(95)
ZET
On the other hand, since ot b-Lipschitz, we have for all x E T:
-v" (dy 0² (x),
which implies
-bλ (v) S inf [bdT v) 2) (x)) ] = (y)
(96)
#ET
Combining both bounds, we conclude that
e²x -b>
(97)
Using this identity, we now bound the dual objective for any (u. v) E K::
f u(xx) u(dx) u(())
/ u(x) v(z))((r)) v(dr)
/ 0' (z)(z)
0 (0(z)
J
I (c)
f 1' v'(x) - bAv(T)
-bA((T) + / (0(z) (du dv).
(98)
Since o² L'' as shown earlier, conclude:
[ (((s) u(dx) / v(x) o(z)) v(dx) S -b) sup / f(x) (dp du).
(99)
jel/
Using the variational characterization of (T, v) (proved earlier), we deduce the upper bound:
ETx (T) + sup seL' / [(x)) (d4 (du du).
(100)
To prove the reverse inequality, let LL and define:
u:=f f. v:= -bb
Then:
u(x) S w(x), v(x) < -b - (-bA - w(xx)) = w(x))
and
u(x) -b- - f(x) -b w(x)
2 -b + inf [bd7(r, - w())
(101)
PET
Moreover, the b-Lipschitz property of yields:
u(x) u(y) = f(x) f(v) - bb < (dr v) X),
(102)
which confirms that (u, v) E K. Applying the variational formula ETx, we obtain:
-bXu(T) 1 f(x) (du dv) / u(x) m(dx) + / v(x) - (T)
(103)
29

---

Since this holds for all f E L', we deduce:
/ f(x) (dp (dp dv) s (T)(H,))
-bxv(T) + sup Jet/
(104)
Putting both directions together, we conclude:
/ f(z) (du dv).
(T, -bAv(T) + sup JEL'
(105)
To recover the symmetric form in Theorem B.2, let f f Furthermore:
bd 2
Then, f L' if and only f E L.
(1(0) (dp du) = / (f(t) 2 2 (dp - dv) = f f(x) (du du)
bd
2 [u(T) - v(T).
(106)
Substituting into Equation (105), we obtain the final expression:
(Tx(2,)) sup SEL / f(x) (du dv) 2 (u(T) + (T)),
bd
(107)
which completes the proof.
D.3 Proof for Proposition B.3
To ensure completeness, we provide full derivations of the result, closely following the methodology
of [42].
Proof. We begin by expanding the definition of the regularized entropy transport:
by
ET ETX = (u(T) + v(7)
+ sup {s (.(M(T)) - x(T] SE
bd w(rr + a. w(r)
bd
a
2
2
{ ( (1 (lo)
+ sup
(# - (dxx): I19ll>>(T) S b
(108)
(na)
We now evaluate each supremum separately: The first supremum corresponds to maximizing a linear function over a symmetric interval. Therefore, evaluates to
by a (u(TT) ~T()).
w(r)
(109)
The second supremum is equivalent to the dual representation of Lipschitz-type transport energy over tree-structured domains. As established in [21, pp. 575-576), we have:
g(wو w(ds)) sup (1 (ld)wdd)))) (u - (w-r)(dr)) (dxx): l l s /
- (V(((a)) -(((()))(((((z))) w(dr).
(110)
Combining both components, we obtain the closed-form expression:
(T = / I(((2))) v(A(x)) w(dr) 2
6A
(u(T) + x(T)
+ w(r) + DA - a (I((T) v(T))
(111)
2
This concludes the proof.
30

---

D.4 Proof for Proposition B.4
To ensure completeness, we provide full derivations of the result, closely following the methodology of [42].
v). This follows directly from the inclusion IL Lo and the dual representation of ETx established in Theorem B.2.
Proof. We begin with the upper bound (Tx ETx (M,
Next, consider a satisfying
bλ
2bL(T) S a -
+ w(r).
(112)
2
We will show that under this condition, the inclusion IL. C LL holds. Then, by Theorem B.2, follows that
( v) & ET)(UD))
(113)
To prove La C. L, we need to show that any function f E La satisfies
bd = f(x) $ w(x) + bd
w(x)
Ve ET.
(114)
2
Let f E Leg. Then by definition,
f(x) =s / ||z]
w(dy)
(115)
where [-w(r) AA + a, w(r) bd -a and |1 L S b. Using this, bound [(x) from above:
f(x) $ 8 + |69| (T) w((r
sw(r)= +
a bL(
2
by
Sw(x) +
a 2bI(T)
2
bd
S <w(x) +
(116)
2
For the lower bound, we have:
f(x) 2 169l (T) (((rxz))
by
2 (r(r
+ a bL(T)
2
bd
2 w(()) (x)
+ a 2bL(T)
2
2 w(x)
(117)
2
Hence, f satisfies the defining constraints of and we conclude that f L. Therefore, La C for all 2 2TL(T).
follows from Theorem B.2 and the definition of ETx that
ETx(M < ET ())
(118)
This concludes the proof.
L
D.5 Proof for Proposition B.5
To ensure completeness, we provide full derivations of the result, closely following the methodology of [42].
31

---

Proof. We first observe that the metric de depends solely on the value of the weight function at the root r of the tree T. This follows directly from the definition of L., where only w(r) appears explicitly.
By construction, have the variational characterization:
da (e, v) sup { 5 (du du) f {EL.
(119)
Let us now verify the metric properties:
(Non-negativity) Clearly, di, (1,, v) 2 0 from the supremum structure. Moreover, dhe (1, μ1) = 0 for
all by linearity of the integral. Now suppose that do (1, v) 0. Using the closed-form expression for ETy, in Proposition B.3, this implies:
(()) + 2 -g]-luul (u((T) u(T]+ S II(A(EE) v(A(x)) w(dr) O.
bd
(120)
Since the first term has strictly positive coefficient by assumption (a < w(r) B)), we must have u(T)= v(7) and
m(A(z)) = (A(()) for all x E T.
(121)
By [42, Lemma A.2], this implies that μ v, establishing identity of indiscernibles.
(Symmetry) Note that if f E La, then -f E La by the symmetric definition of the function class. Therefore, from Equation (119), obtain
() = du(,,)).
(122)
(Triangle Inequality) The triangle inequality holds immediately from the supremum definition over convex, symmetric function class:
da(u,o) + da(a,v) 2 (dp- do) (
/ 5 (do dv) / [dp (dp du).
(123)
for all f E Log, and taking the supremum yields the inequality.
Hence, du satisfies all properties of metric on .M(),, and the proof is complete.
D.6 Proof for Equation (13)
Proof. We recall Equation (13). Let f E (-R(IR)) be non-negative density function. The Radon Transform R maps / to density defined on a tree system T, while preserving the total mass:
Il = [ R7flI,. 1
for all TE
(124)
To establish this property, first observe that the non-negativity of ensures that the transform preserves non-tegativity: if f 2 O, then R 2 O, implying that the transformed function is a valid density. The preservation of total mass then follows directly from the definition of Ro, which integrates over linearly parameterized subsets aligned with the structure of T.
|R7
Σ 1ET
/ R f(z, 1)1 dt.
Σ / 1 f(y) aly, C): (t, - -(y 21, 01) dy dt,
Σ / ( f(v) alv, C), (tx - (y x1, 01) dy dt
Tec JR
El (f((0) aly,
- C): -8 (tx (y xz 01)) dtz dy
lEL
32

---

Σ J f(y) aly, (Sa 8 (tx
(y 21, =1, 01)) dtz) dy
LEC
{ f(y) aly, L) dy
Σ
lEC
JSa f(y) "Day L) dy
IEL
J f(v) dy
11f1[1..
(125)
The proof is completed.
D.7 Proof for Theorem 3.3
Proof. We consider the expression
PartialTSW(j,.) = S1 da daTT, UT) do(
(126)
and show that it defines metric on M(R"). Since the splitting map is E(d) -ivvrinnn the Radon Transform R is injective; that is, for any f E I.- (TR)), if Rg = 0 for all E T, then f (see [80]). We now verify the three properties required for PSr to be a metric on M(R)).
Positive definiteness. For μ. E P(R"), it is clear that PartialTSW((,, p)) PartialTSW (14, v) 2 0. Morcover, if PartialTSW(p, v/) = 0, then du(HT, VT) = 0 for all T E T. Since da is metric on MMT), it follows that for VT for all T. Hence, R4 R4 for all T E T. By the injectivity of R2, we conclude that fn = fu. and thus μ V.
0 and
Symmetry. For any J, M(").. we have:
PartialTSW(H4p)) v)
da daT do
J, daluT HT) do(T)
(127)
PartialTSW(y.M)
(128)
Therefore, PartialTSW(I1,) v) = (Mr
Triangle inequality.
For (1, M2, (RR)), compute: PartialSS((11,P))) PartialSSW(H2,
da da(Hi.T,P2 do + da(H2.. dT
[ (du(H)., da dT(T)
de(H..T do
PartialTSS (Prta
(129)
where the inequality follows from the triangle inequality satisfied by d, on each tree T. conclusion, PartialTSW satisfies all properties metric on the space M(Rd).
We aim to show that PartialTSW is E(d)-invariant, meaning that for any E(d), the following holds:
PartialTSW(H, = PariaaTSS(Ggha.grr
(130)
33

---

where Gü and gtv denote the pushforwards of and respectively, under the Euclidean transforma- tion g: R2 R
Let T E T be tree system given by T = {4 = (2,, 0,)) Then, under the action of = (Q,a),
have
gT {gl, (Qx; a, Q0:)
(131)
We also note that gifn = foin and gifu = faiv. Since det(()) = 1, we compute:
R Rgclgdf))or
14 (gtf aly, 9C), (tgz - (y - Tgl. 0g1)) dy
J Jug v) ((y.ggL) 5 81 fg 0gt)) dy
[ fulo 98) algy. gC)1 ō (tz [gy - gt 0gt) d(gy)
[ ful aly. L)i (1. (gy at fg)) dy
J fu()) aly, C)i 8 (t. (Qy + Qzi a, Q0i)) dy
344 fu(U) ay C) 8 (tx (Q .1), Q81) dy
fu(v) aly, C)1 8 (Ia (y 2Ι, 01)) dy
(132)
similar computation gives:
Rye (gc(ot(v)ggn gl) = R2,x(x,1).
(133)
Moreover, since g acts isometrically tree systems, the induced measures satisfy:
du da(pr, = du ((ott))), (ott)
(134)
Thus, we compute:
Partia/SS(sttu.ger
da (((ott)) do
da ((9i)))) gT, (9th) do (gT)
S. du (UTT UT) do(g))
S da da VT) do
(p.,
(135)
We conclude that PartialTSW E(d)-invariant.
Remark D.1. We omit almost-sure conditions in the above proof, as they are straightforward to verify and would otherwise obscure the main argument.
E
Experimental Details
E.1 Algorithm for Partial Tree-Sliced Wasserstein Distance
The computation of the Partial Tree-Sliced Wasserstein (PartialTSW)) distance is outlined in Algo- rithm This procedure estimates the distance by averaging costs derived from multiple tree-based projections of the input measures.
34

---

Algorithm 1 Partial Tree-Sliced Wasserstein distance.
Input: Measures j and in M number of tree systems I., number of lines in tree system space of tree systems T, splitting maps o, parameters a, b, X, total mass ((T), (T)..
Scale total mass of fa and 1/ such that (T(R)) = (TT), 2 [R)) = =(T).
for i = 1 to L do
Sampling x E TRd and 0 D....., Ok lid U($d0-2))
Contruct tree system L; = ((r, 01) fi.,................. (xr, }-
Projecting μ and v onto T/ to get R2, 1 and R2, V.
Compute PartialTSW(p,, (p, .) (1 da (R22, Hi R2, v).
end for Return: PartialTSW (1,, v).
E.2
Computational and Memory Complexity Analysis
This section details the computational and memory demands of our proposed PartialTSW distance. We consider input measures and represented by N samples in a d-dimensional space, with L tree constructions and & lines per tree.
Table outlines the complexity of key operations. The dominant factors are the distance- based weight splitting (OL)Nd)) for projecting samples and the sorting of these projected ID coordinates (OLLN log N)). Consequently, the total computational complexity is O(LNNd + L&N log N). The primary memory consumers are the storage of split weights, tree/line parameters, and the original data, leading to an overall memory requirement of O OLLN + Lkd + Nd).
Table 3: Detailed complexity analysis for PartialTSW. (N number of samples, = dimension, L number of trees, k = lines per tree).
Operation Category
Specific Steps Involved
Computational Cost
Memory Cost
Adjusting sample weights for a and to meet target total masses.
O(N)
O(N)
Initial Mass Scaling
Distance-Basse Weight Splitting
Calculation of distances from N O(LANd) points to Lk lines, and subsequent softmax for weight distribution.
O(LAN + Lkd + Nd)
Sorting Projected Data
Sorting the N projected coordinates O(LAN log N) along each of the Lk lines.
O(L&N)
Overall Total
O(LNNd + Nk
O(LEN + Lkd + Nd)
GPU Memory Optimization for Distance- Based Splitting. The practical GPU memory footprint for the distance-based splitting step can be significantly lower than a naive theoretical estimate. As highlighted by [80], this operation involves (1) computing d-dimensional distance vectors from points to lines, (2) calculating their norms, and (3) applying a softmax function across lines within each tree obtain split weights. While direct implementation might suggest O(LkNd) memory for storing all intermediate distance vectors, modern deep learning frameworks like Py when using compilation tools (e.g., forch.compile. sequential computations into fewer GPU kernels, potentially allowing large intermediate tensors (like the full L&N x d distance vectors) reside in faster, smaller shared memory or be recomputed on-the-fly, rather than occupying global GPU memory. Consequently, the persistent global memory primarily stores the essential data: line parameters (O(Lkd)), sample coordinates (O(N)), and the resulting split weights (O(LN))), aligning with the O(LNN + Lkd + Nd) overall memory profile.
can perform kernel fusion. This optimization merges these
E.3 Empirical Runtime and Memory Performance of PartialTSW
We present an empirical evaluation of the runtime and memory usage of PartialTSW. The experiments were conducted on a single NVIDIA H100 GPU. We fixed the number of tree iterations L, = 10 and lines per tree k 4. The analysis varies the number of samples N E {100, K, 5K, 10K, 500K) and the data dimension d (50, 100, 500, 10003.
35

---

30
50
10000
50
25
Memory usage (M1B)
8000
Runtime (mc)
8000
1.5
2000
40000
Number supports
Number of supports
Figure 7: Empirical runtime (left) and peak memory usage (right) for PartialTSW, varying the number of samples (N) and data dimension (d). (L = 10, 4).
Runtime Scalability. The empirical results, depicted in Figure (left), illustrate how the runtime of P scales with the number of samples N and the data dimension d. The runtime exhibits near-linear increase with N. For instance, processing N 50, 000 samples takes approximately five times longer than = 10, 000 samples (when d, I,, k are fixed), which is consistent with N from our theoretical analysis (Section E.2). Regarding
the O(Nd + log .) dependency dimensionality, the runtime also demonstrates linear dependency on d. For example, increasing d from 10000 50000 5x increase) results in correspondingly proportional increase in runtime for fixed N. This aligns with the (d) factor in the Lk:Vo d term of the complexity. These empirical observations support the theoretical computational complexity.
Memory Scalability. Figure 7 (right) showcases the memory consumption characteristics of Par- tialTSW. The peak memory usage scales linearly with both the number of samples N and the dimension d. This behavior is predictable and directly corresponds to our theoretical memory com- plexity of O PLLN + Lkd Nd), indicating efficient memory utilization that grows manageably with data size and dimensionality.
E.4 Sample Complexity and Estimator Stability
We first clarify the roles of the two key parameters in method.
Lines per Tree (k)): This is a structural parameter that defines the ground-truth distance PartialTSW&& v). As discussed in the TSW literature, using k: > enhances the method's capacity to capture complex topological and structural features of the data.
Number of Trees (L): This the Monte Carlo (MC) estimation parameter. Our method approximates the ground-trat truth distance by averaging over L independently sampled random trees.
As discussed in the computational complexity analysis in Section 3.3, the total computational cost is proportional to the total number of 1D projections, which is = L.k. For a fixed computational budget N, this creates natural trade-off: increasing improves topological expressiveness, but requires decreasing L,, which in turn affects the stability the MC estimate.
In experiments, we tune the structural parameter k for empirical performance and then set L= N/k to ensure the total number projections N remains fixed, allowing for fair comparison against baselines under the same computational budget.
This approach raises a valid concern about the stability of the estimator, as the MC approximation for fixed & decreases the estimator's convergence. We measure stability using the Coefficient of Variation (CoV = a pt), a normalized metric where is the standard deviation and μ the mean of the distance estimate over multiple runs.
rate of (II-1/2). We therefore provide empirical analysis of
Table 4 shows the CoV function of both the MC parameter and the structural parameter k. The results empirically verify the expected convergence. For any fixed k the estimator's stability improves (i.e., CoV decreases) as the number of MC samples increases, aligning with the theoretical (L2-1/2) rate. For instance, at & 5, increasing L from 10 to 1000 reduces the CoV by over 14x.
36

---

Table Estimator stability analysis. The table shows the Coefficient of Variation (CoV o/p) for the PartialTSW distance as function of the number of tree slices (L) and the number of lines per tree (k).
Number of Trees (L)
k
k 10
k
100
10
0.1098
0.1016
0.0456
50
0,0526
0,0252
0.0183
100
0.0263
0.0239
0.0091
500
0.0105
0.0133
0.0070
1000
0.0076
0.0070
0.0041
10000
0,0029
0.0018
0.0023
20000
0.0016
0.0024
0.0010
This analysis confirms that in our large-scale experiments, where use high number of projections (e.g., L 2 1000), the resulting distance estimate is stable and reliable.
E,5
Discussion on Hyperparameters of Evaluated Methods
This section briefly outlines the key hyper-parameters for each evaluated Unbalanced Optimal Transport (UOT) and Partial Optimal Transport (POT) method and their respective roles.
SPOT 19]. The hyperparameter & specifies the number of points to be transported, thereby defining the partial nature the matching between distributions.
SOPT [2]. The regularization parameter controls the "partialness" of the transport by influencing the total amount of mass that is optimally transported between distributions.
Sinkhorn [74]. The hyperparameter reg the entropic regularization coefficient that smooths the optimal transport plan. The hyperparameter regm is the marginal regularization coefficient that penalizes deviations from the prescribed marginal constraints, thus allowing for mass variation.
SUOT and USOT [7]. The hyper-parameteer P1 and P2 are regularization parameters. They respectively control the cost of deviating from the source and target marginals in the sliced domain, enabling unbalanced transport by permitting mass creation or destruction.
PAWL [13]. The hyperparameter the number of points to be transported, effectively determining the extent of partiality in this unbalanced optimal transport formulation.
UOT-FM [22]. The hyperparameter À influences the regularization of marginal constraints, thereby controlling the degree which the masses of the coupled distributions must be preserved during transport.
ULightOT [28]. The hyperparameter governs the extent of mass conservation, adjusting how strictly the total mass of the transported distribution must adhere to the original or target masses.
Partial-TSW (Ours). The mass parameter v(T) specifies the proportion of the target distribution's 's mass to be matched by the transport plan. The source distribution's s mass proportion, μ(T), is typically fixed at 1, adjusting ((T) controls the partiality of the matching against the target.
E.6 Comparing Computational Efficiency
To ensure consistent and fair results, two warm-up runs were performed for each method and each sample size n before conducting 10 timed repetitions. The average runtime and peak memory usage (for GPU methods) were then recorded. Unless otherwise specified (as the discussion on varying d below), these experiments were conducted with data of dimension d 2 and for sample sizes ranging from 102 to 10
Since hyperparameter choices can significantly affect algorithmic runtime, the specific settings used for each method in this runtime comparison are detailed below. For general description these hyprpaarametess and their roles, please refer to Appendix $E.5.
37

---

SOPT SPOT
Sinkhom USOT
PartialTSW
PAWL
Runtime (s)
128
1024 2048
Dimension d
Figure 8: Runtime comparison for PartialTSW and POT/UOT solvers data dimension d.
Common settings for the compared sliced-based methods (SOPT, SPOT, USOT, SUOT, PAWL)
included L 10 projections. For PartialTSW (Ours), we used num_trees = 5 and num num_linee lines This configuration for PartialTSW, where the product of num_trees a comparable number of one-dimensional sorting operations to the = 10 setting in other sliced methods, aiming for fair comparison. Specific hyperparameters for each method were then set follows:
num limelines = 10. offers
SOPT: Regularization parameter 1.0.
The number of transported points was (the input sample size for each distribution), implying full matching was performed. (Number of projections L = 10, as stated above).
SPOT:
Sinkhorn:
Entropic regularization reg 0.1, marginal KL regularization regm 1.0. maximum number of Sinkhorn iterations 'numltermax" 100, and stopping threshold 'stopThr' =0-
USOT and SUOT: Regularization parameters P1 0.01 and P2 1.0.
PAWL: The number of transported points was set (implying full matching).
PartialTSW (Ours): The target mass proportion v(T) was set to 1.0 (with the source mass proportion ju(T) typically assumed to be 1.0). This choice was made because v(T) does not affect the computational runtime of the PartialTSW implementation used this benchmark.
Furthermore, present runtime comparison for varying data dimensions d in Figure 8, The results indicate that the runtime is not significantly affected when d increases.
The runtime comparisons for all methods were conducted with an Intel Xeon Platinum 8580 CPU and an NVIDIA H100 GPU.
E.7
Noisy Point Cloud Gradient Flow
We used clean point cloud data obtained from [2] for the dragon and bunny shapes. Each clean dataset contains 10k data points. We randomly select and add 7% noise points to the target point cloud (bunny). Inspired by [2], the noise is sampled from the region [-0.6M, 0.6./ where M max4e 1.n) (I]=,I)), where z; is the point in the target. In total, the target point cloud consists of 10k clean points and an additional 700 noise points. We use = 10 projections for SW, and L= 5 trees, l = 2 lines for TSW and PartialTSW. All methods trained using Adam optimizer with learning rate 10-3 over 300 epochs. The results are shown in Figure 3.
All experiments were conducted with an Intel Xcon Platinum 8580 CPU and NVIDIA H100 GPU.
38

---

E.8 Robust Generative Model
E.8.1 Implementation detail
Pre-training an Autoencoder (AE). An Autoencoder (AE) is pre-trained provide 2D latent representations E R2 for MNIST digits. We employ Wasserstein Autoencoder with MMD regularization (WAE-MMD) [77] architecture. The AE is trained for 50 epochs using the Adam optimizer with learning rate of 3 x 10 and a batch size of 256. The WAE-MMD loss uses a À hyperparameter of 500.0 to balance reconstruction and MMD regularization terms. For the MMD term, we match the aggregated posterior g(=) to uniform prior distribution p(z) ~ U(-1, 1J2. This encourages the learned latent space to reside approximately within [-1, 1j The training data for the AE consists of MNIST digits 0 and balanced and augmented as described below. The latent dimension is set to =
The Autoencoder, AE [0, 1|111x2222 [0, 1_1_xxxxxx architecture is as follows:
Encoder:
Input: 1 x 28 x 28 (MNIST image)
channels 2. padding = 1) - ReLU (32 x 14 14)
out cut_channels
32, kernel_sizee
4, stride
(32, 64, kerneeeeeiiiiie size = 4, stride = padding = - ReLU
(64 x 7
Flatten: 64 x 7 x 7 = 3136 features
Linear (in features = 3136, out features = 512) - ReLU
Linear 512, latent_di dimm_dimm = 2) (for mean j1)
Linear (512, latent_ dim = 2) (for log-variance log o²)
Latent vector μr + (Reparamettrizzation trick)
Decoder:
Input: 2 E I gllteet_diig dim
Linear Linear(latent dim = 2, 512) - ReLU
Sinear(512, 3136) - ReLU
Reshape to 64 x 7 x 7
ConvTransposee2 (64, 32, kemee_size 4, stride
2, padding
1) ReLU
(32 x 14 14)
GoovTransposee2 1, kernel size = 4, stride = 2, padding = 1) Sigmoid (1 28 x 28)
Dataset Augmentation for Auxiliary Models. To ensure robust training of the AE and the digit classifier, we prepare balanced and augmented training set from MNIST digits and 1. The original MNIST training set contains an unequal number samples for these digits. We balance these classes by applying data augmentation
the minority class until its sample count matches the majority class. Augmentations include random affine transformations (degrees: 115%, translation: =0.15 image dimension, scale: 0.85 1.15x) x and random rotations (+15"). This balanced and augmented dataset used exclusively for training the AE and the binary (0 vs. 1) digit classifier. We found that having balanced dataset for training AE would lead balanced latent space for MNIST Digit and 1.
Pre-training an MNIST Digit Classifier. convolutional neural network classifier is pre-trained to distinguish between MNIST digits 0 and 1. It is trained for 20 epochs on the balanced and augmented dataset of these two digits, using the Adam optimizer with a learning rate of 1 x 10-3 and Cross-Entropy loss function. This classifier achieves approximately 99, 99.99 accuracy set of unseen MNIST Os and and is subsequently used (with frozen weights) to evaluate the class labels of images generated by the main generative model.
test
The Classifier, C: [0, 1] x28x28
R architecture as follows:
Input: x 28 x 28 (decoded image)
Conv2d(1, 32, kernel__iiiiiee size = 3, stride 1, padding = 1) - ReLU (32 x 28 x 28)
MaxPoo12d(kernel_ size stride 2) (32 x 14 14)
Conv2d(32, 64, kernel size 3, stride 1, padding 1) ReLU (64 X 14 X 14)
39

---

MaxPoo12d(kernel size = stride 2) x
Flatten: 64 x 7 x = 3136 features
L31ear 128) ReLU
Linear(128, wwmnnnummaaaaaaa classes = (Logits for classes 0 and
num_classes
Constructing the Observed (Contaminated) Dataset X The observed dataset X for training the generator G consists of latent representations. These obtained by encoding MNIST images of digits 0 (target class) and (outlier class) using the pre-trained AE's encoder. Specifically, Xobs mixture comprising 90% samples from the true latent distribution of digit O (X)) and 10% samples (outliers) from the true latent distribution of digit 1 (21). To construct this, we sample latent vectors z' from the prior [[- (1, 1]2, decode them images x'' A (='), and classify x' using the pre-trained O/1 classifier. If x is classified as O (or 1), 2 is added to pool for Xo (or X1). We collect samples until we can form a dataset of Nobs 50, 000 latent points, with the 90/10 proportion.
These latent points constitute Xobs and are scaled to approximately reside within [-1,
Training the Generator G. The generator G (2) T2) [-1, 1]2 is multi-layer perceptron (MLP) designed to map 2D Gaussian noise Z ~ N (O I2) to the target latent space. The generator is trained by minimizing (Partial) Optimal Transport distance D(G(Z). Xots), where Z is batch of noise samples. Training is performed for 30 epochs using the Adam optimizer with learning rate 2f 10 and batch size of 256. Specific (P)OT-based distances D used for PartialTSW and baseline methods are detailed in the main paper.
The generator architecture is:
Input: Z E R2 ~ 12)
Linear(2, 4) BatchNormid(41 Lea)yReLU(O.2)
Linear( 8) BatchNormid(() LeakyReLU(0.2)
Linear (8, 2) Tanh (Output Zgen E 1|²)
Evaluation. To evaluate the generator's ability to learn the target distribution Xo while ignoring outliers from X1, we employ two main criteria:
1. Outlier Rate:
000 latent samples Zgen = G(2). These latent samples are decoded into images å = A Faee (Zgen) using the pre-trained AE's decoder. The resulting images are then classified by the pre-trained 0/1 digit classifier. The outlier rate is the percentage of generated images classified as digit 1. A lower rate indicates better robustness.
We generate Neval
2. Sample Quality and Diversity: We qualitatively assess the generated samples by visualizing the decoded images and their corresponding latent representations āgen- We look for high-fidelity generation of digit and good coverage of its variations, indicated by well-distributed latent space for the generated samples classified as 0.
Performance summaries, including outlier rates and visual comparisons, provided in Figure 4 and Table 1 in the main text.
Hardware Settings. The experiments for all methods were conducted on system equipped with an Intel Xcon Platinum 8580 CPU and one NVIDIA H100 GPU.
E.8.2 Ablation result for baselines
We evaluate the impact of hyperparameter settings each method's ability to isolate the target MNIST distribution from the 10% MNIST outliers present the training data. The following summarizes these ablation results (Figures 9-15), focusing the percentage of generated MNIST outliers and, where applicable, qualitative aspects of the learned distributions and generated samples.
SPOT [9]. Figure 9 demonstrates SPOT's varying success in isolating the target MNIST data from the 10% 1 outliers, contingent on its hyperparameter yielding 45.62% MNIST 1 outliers) or very large l values (e.g., & = 256, yielding 16 16.20% outliers) result poor outlier rejection, an optimal range for i around 200 210 reduces the MNIST outlier rate to 6 - 7%. This indicates substantial but incomplete removal of the 10% outliers.
While very small values (e.g., k 10,
40

---

SOPT [2]. SOPT's effectiveness in discarding the 10% MNIST outliers is modulated by its regularization parameter A, as shown in Figure 10. The lowest outlier percentage achieved by SOPT is 13.28% (at À 0.01), which still exceeds the initial 10% contamination level. Larger values of À lead to even higher and relatively stable outlier rates (around 15 16.42%), indicating persistent difficulty for SOPT in cleanly separating the target distribution in this setup.
Sinkhorn [74]. Sinkhorn shows the potential for complete removal of the 10% MNIST 1 outliers when its entropic regularization reg and marginal regularization regm are appropriately co-tuned (Figure 11). Specifically, setting reg = regm values of 0.5, 0.7, or 0.9 results in 0% MNIST outlier generation, successfully achieving the task's s objective. However, imbalanced or overly small regularization values lead substantial outlier contamination (e.g.. 52.48% for reg regm 0.3, or 70.56% for reg 0.9, regm = 0.1). Moreover, qualitative inspection of the results (Figure 11, particularly for reg = regm (0.5, 0.7, 0.9)) reveals that while Sinkhorn effectively removes outliers, the generated latent distribution for MNIST O digits appears clustered, and the corresponding decoded images may lack diversity compared to the true distribution. This suggests a potential trade-off between perfect outlier rejection and capturing the full diversity of the target class for this method under these settings.
SUOT [7]. The performance of SUOT in the task of removing 10% MNIST outliers is consistently poor across the explored range of its marginal regularization parameters f1 and p2. as detailed in Figure 12. The method yields a high MNIST outlier rate of approximately 41% regardless of the hyperparameter settings tested. indicating a failure to distinguish the target MNIST distribution from the contaminants.
USOT [7]. USOT, while performing better than SUOT, still struggles to fully reject the 10% MNIST 1 outliers (Figure 13). Across the tested range of its #1 and p2 hyperparameters, USOT yields consistent MNIST outlier rate of approximately 17.08%. This suggests that while mitigates some contamination, it does not fully isolate the target MNIST 0 distribution in this scenario.
PAWL [13]. PAWL demonstrates exceptional success in the goal of removing 10% MNIST 1 outliers, as shown in its ablation study (Figure 14). It consistently achieves a 0% MNIST 1 outlier rate across all tested values of its hyperparameter is (from 10 to 256). This indicates PAWL's strong capability to identify and learn the target MNIST O distribution while completely ignoring outliers, exhibiting robust performance across wide range of k. However, noted in the main text and suggested by qualitative inspection of Figure 14, this strong outlier rejection by PAWL may be accompanied by less sample diversity, with its learned latent space showing heavily concentrated clusters.
PartialTSW (Ours). Our PartialTSW method shows strong capabilities in removing the 10% MNIST 1 outliers, with performance critically depending on its mass parameter n(T) (Figure 15). Complete outlier rejection (0% MNIST 1 outliers) is achieved for (T) values between 0.3 and 0.6. Setting v(T) closer to the true inlier fraction of 0.9 (which yields 9.72% MNIST outliers) leads to the model fitting the outliers. This highlights that optimal robustness for PartialTSW is achieved when D(T) is chosen to be somewhat less than the actual inlier data proportion in the contaminated dataset. Qualitatively, as seen in Figure 15 and highlighted in our main findings, the settings achieving complete outlier rejection (e.g., ((T) E [0.3, 0.6)) also yield well-distributed latent space and diverse image samples for the MNIST O class, effectively capturing the target distribution.
41

---

MNIST
MNIST
0
0
0
0
0
0
0
0
0
0
0
(a) Dataset with 90% MNIST 0 and 10% MNIST
(b) = 10
45.62% MNIST
0
0
0
0
0
1
0
0
0
0
0
0
0
0
0
1
0
0
0
0
0
0
0
1 0
1
0
1
1
0
0
0
0
1
0
0
0
0
0
0
0
0
0
0
0
0
0 0
0
0
0
0
0
0
1
0
/
1
0
0
0
0
0
0
0
0
0
0
0
0
1
1
1
1
1
0
0
1 0
0
0
|
0
1
0
0
0
/
0
0
0
0
0
1
0
0
0
0
0
0
(c) k 25
16.36% MNIST
(d) 50
27.58% MNIST 1
0
0
0
0
0
0
0
0
0
0
0
1
0
1
0
0
0
0
0
0
0
0
0
0
0
0 1
0
0
0
0
0
0 0
0
0
1
0 0
0 0
0
0
0
0
0
(c) k 75
18.14% MNIST
kf) 100
14.86% MNIST
0
O
0
0
O
(g) k: = 200
6.32% MNIST
(b) k: = 200
7.12% MNIST 1
O 0
O
O 0 O 0 0
0 0 /
O 0 0
0 0 0
O
0 0
0
0
(1) 210 + 7.06% MNIST 1
(i) A = 256
16. 20% MNIST
Figure Ablation study of SPOT [9] for robust image generation. The figure illustrates the percentage of generated MNIST 1 digits (outliers), along with the corresponding learned latent distributions (left) and decoded image samples (right).
42

---

MNIST
MNIST
(a) Dataset with 90% MNIST O and 10% MNIST
(b) A 0.01
13.28% MNIST 1
0
/ O O
O O 0
O O
O
(c) A = 0.05 -
15.00 02%
MNIST
(d) X = 0.1
16. 20% MNIST
0
0
0
O /
O
O / 0
0
0 1
0
/ 0
0
(e) = 0.5- 16.12% MNIST
00
10.20% MNIST
O
/ 0
O I
O
O 0
0 0 0
/ 0
0 0
x = 10
16.20% MNIST
(h) A 100
16.20% MNIST
Figure 10: Ablation study of SOPT [2] for robust image generation. The figure illustrates the percentage of generated MNIST digits (outliers), along with the corresponding learned latent distributions (left) and decoded image samples (right).
43

---

MNIST
MNIST
0000100
0 00000000
0
0
0 00001
1
0
0 0
100
0
(a) Dataset with 90% MNIST O and 10% MNIST
(b) reg = 0.3, regm = 0.3 - 52.48% MNIST
0
0 O
O
0
0
0 0
0
(c) reg- = 0.4, regm 0.4 23.76% MNIST
(d) reg 0.5, regm
0.5 0% MNIST
(c) reg regm
0.7 +
0%
MNIST 1
(f) reg 0.9, reg..
0.9
0% MNIST
0
0
1
0
0
1
1
0
1
1
1
0
101
000
1
1
11000
1
0 010100
10000
0 0
0
(g) reg- = 0.9, regm-
0.1 - 70.56% MNIST 1
(h) reg regm = 0.9
14. 1..1%% MNIST
0 0
0 1 0
0
0
1 0
0
0
0 0 00
0 0
110011
0 1 0100001100
0
0 1 0 1001000
1 0 0
0
0
0
0
(i) reg = regm. 0.3 + 40.00% MNIST
( reg = 0.3, regm = 0.5 - 14.92% MNIST 1
Figure 11: Ablation study of Sinkhorn [74] for robust image generation. The figure illustrates the percentage of generated MNIST digits (outliers), along with the corresponding learned latent distributions (left) and decoded image samples (right).
44

---

MNIIT
O
0
0 0
/
O
0
0000
00
100000
000000
0
0
0
0
(a) Dataset with 90% MNIST 0 and 10% MNIST
(6) P1
0.01 P2 0.01
41% MNIST
0
0
0
0000
0000000
0
0
0
0
/
O
1
0
0
0
/
\
1
0
/ /
0
1
0
/
O / 0 1000
0
1 1000
0 / 00000000
OO
0/00000
00000011
0 0 / 0 10/1000
0 0
0 0
0
0
/
0
O
O
0
0
0
(c) 1 =
41% MNIST
(d) = 100,
100 - 41% MNIST
0
O
O
0
O
0
0
0
0
0
/
O
1
O
0
0
0
1
1 1
/
0 0
0
0 I / / 1000100
0
0
0
O / / / 0 1000
O / 10000000
0O / 00000000
/ 001 000000
00000000
0 / 1 1/1000
0 / 10/100 00/1000
0
/ / 0
10011
0 0
/ 0 00002
/
0
/
0
0
0
0
O
0
0
0
0
0
0
0 0
0
0
0
0
(e) P1
200, P2 200
11% MNIST
( P1 1000, P2 1000
41%
MNIST
0
0
1
0
/ /
0
00000
0
0
O
0
10000
0
0
0
0
0
/
1 0
O
/
O
0
0
0
1
1 1
100
0
\
1
/
0
0
00
1010/1
0
0
1
/
O
/
/ 0 000/0000 1000
O
1
0
/ 0000
/
0
/ 10000000
0
0
0
0
0
/ 0
10/10000
0
1
0
/
1
0
0
/
/
0
0 001
0
/
0
/
0
0
0
0
0
0
0
0
0
0
(g) PI P2 1000
41% MNIST
(h) P1 1000, P2
41% MNIST 1
Figure 12: Ablation study of SUOT [7] for robust image generation. The figure illustrates the percentage of generated MNIST digits (outliers), along with the corresponding learned latent distributions (left) and decoded image samples (right).
45

---

MNIST
MNIST
(a) Dataset with 90% MNIST 0 and 10% MNIST
(6) p1 0.01, P2 0.01
17 17.08%
MNIST
0
( 12 =1
17. .8% MNIST
(d) = 100, P 100
17.08% MNIST
0
0 0 0
0 5
0 0 0
O
0 0
0 0
0
O
(0) P1
200, P2 200
17.08% MNIST
( P1 1000, P2 1000
17.08% MNIST
0
O 0 0
0 1 /
O 0 0 0
0 0 0
O O
0
(g) P2 1000 17.08% MNIST
(b) P 1000, P2
17.08% MNIST
Figure 13: Ablation study of USOT [7] for robust image generation. The figure illustrates the percentage of generated MNIST digits (outliers), along with the corresponding learned latent distributions (left) and decoded image samples (right).
46

---

MNIST
MNIST
(a) Dataset with 90% MNIST 0 and 10% MNIST
(b) k = 10
0% MNIST 1
(c) k = 25
0% MNIST
(d) k = 50
0% MNIST
(c) 75
0% MNIST
(f) k = 100
0% MNIST
(g) k 200
0% MNIST
(h) 256
0% MNIST 1
Figure 14: Ablation study of PAWL [13] for robust image generation. The figure illustrates the percentage of generated MNIST digits (outliers), along with the corresponding learned latent distributions (left) and decoded image samples (right).
47

---

MNIST
MNIST
(a) Dataset with 90% MNIST 0 and 10% MNIST
(T) = 0.3
0% MNIST
(c) v() 0.4
0% MNIST
(T) = 0.5
0% MNIST
(7) 0.6 0% MNIST
(T) = 0.7
0.08% MNIST
(T) 0.8 2.18% MNIST
(b) xT) 0.9 9.72% MNIST 1
Figure 15: Ablation study of Partial-TSV -rsiall-WWWWao (Ours) for robust image generation. The figure illustrates the percentage of generated MNIST digits (outliers), along with the corresponding learned latent distributions (left) and decoded image samples (right).
48

---

E.9 Imbalance Image to Image Translation
E.9.1 Implementation detail
This section outlines the experimental setup for the imbalanced image-to-image translation task, specifically converting "Young" faces to "Adult" faces.
Dataset. Our experimental dataset and preprocessing follow [41]. We utilize the FFHQ dataset [37] of 1024 x 1024 images. These images are encoded into a 512-dimensional latent space using pre-trained ALAE autoencoder [64]. The resulting latent representations are categorized into two classes: "Young" and "Adult", based on a cutoff age of 45 years. This process yields an imbalanced dataset comprising approximately 38,000 38,000 "Young" latent vectors and 10,500 "Adult" latent vectors. The translation is performed within this 512- 12-dimensional latent space.
Translation Accuracy. To evaluate the accuracy of the translation from the "Young" to the "Adult" domain, we adapt the procedure from [28]. A classifier is pre-trained on the 512-dimensional latent vectors to distinguish between "Young" and "Adult" images. This pre-trained classifier, which achieves 99% accuracy on a held-out test set latent vectors, is then used to assess whether the translated latent vectors (X) (where X are latents from the "Young" domain) are correctly classified as "Adult".
Perceptual Similarity. We measure the perceptual similarity between the original images (recon- structed from X) and the translated images (reconstructed from M (X)) using the Learned Perceptual Image Patch Similarity (LPIPS) metric [90]. For LPIPS calculations, with pre-trained weights. While some prior work, such as [28]. employed attribute-specific metrics like "Keep Accuracy" (e.g., for preserving gender), selected LPIPS to offer
use the AlexNet backbone
comprehensive assessment of overall visual fidelity post- cost-translation, rather than focusing on a single attribute.
UOT-FM Baseline. We compare against Unbalanced Optimal Transport Flow Matching (UOT- FM) [22]. Following [28], we parameterize vector field ti using a 2-layer feed-forward network with 512 hidden neurons and ReLU activation. We apply their default configuration for Flow Matching. Consistent with the approach [28], we perform an ablation study over the regularization parameter
X, which controls the penalization of deviations from marginal constraints.
ULightOT Baseline. We also include ULightOT [28] as baseline. We adapted the publicly available code and default models for our experiments. Following the methodology in [28], we ablate the parameter 7, which governs the degree of mass conservation in the transport plan. Our empirical observations indicate that the performance of ULightOT saturates for т 1000. For instance, increasing 7 to 10000 yielded negligible changes the Accuracy-LPIPS trade-off compared to 1000, as shown by the results (e.g.. in Figure 5 of the main text).
Mapping Network Architecture.
For methods such as SW, Db-TSW, and our PartialTSW, the mapping network M is implemented us- ing ResidualMLP. The input to this network is latent vector E 1212 The specific Residuall dualMLP configuration used has an input/output dimension of 512, with num num_hidden_ wet_hidden_blocks=0000 hidden dim siimmmultipli mulmippieerrr....
blocks=0 and
The core of this network, denoted as MLP core, processes the input through the following sequence of operations:
Apply an initial linear transformation: Linear(512,511 512)
Followed by layer normalization: LayerNorm
Then, apply the GELU activation function: GELU()
Apply dropout with a rate of 0.1: (rapout((.1)
Finally, apply an output linear projection: Linear(512,5))
Let the output of this sequential MIP core block be MLP (z).
The final output of the mapping network (2) is obtained by adding a scaled residual connection to the original input:
M (z) + MLP (2)
where a is learnable scalar parameter (analogous to LayerScale) that is initialized to 0.1.
49

---

Table 5: Comparison of model size and training time per epoch for the Young-to to-Adult translation.
Method
Number of Parameters
Time (s/epoch)
SW
526,337
Db-TSW
526,337
25
PartialTSW (Ours)
526,337
25
UOT-FM
788,224
35
ULighttT
5,263,380
60
Table 6: Full image-to-image translation results, averaged over runs. Our method, PartialTSW, consistently demonstrates superior trade-off between Accuracy (T) and LPIPS ()) across all translation directions.
W-M
M--W
A-Y
Method
Parameter
Ace (%)
LPIPS
Ace (%)
LPIPS
Acc (%)
LPIPS
SW
93.95
0.4418
91.68
0.4546
89.82
0.4041
D--TSW
94.13
0.4436
92.07
0.4546
89.58
0,4022
0.0005
49.21
0.3914
85.52
0,4269
79.77
0.3836
0.005
70.69
0,4531
94.19
0.4749
92.91
0.4129
UOT-FM
0.05
80.28
0.4899
95.40
0.5106
97.40
0.4693
81.50
0.5198
97.91
0.5369
98.43
0.4828
50.0
76.36
0,4102
85.91
0.4086
84.31
0.3452
250.0
86.36
0.4466
92.81
0.4516
90.49
0.3906
ULightOT
1000.0
88.07
0.4557
93.91
0.4626
92.00
0,4060
10000.0
88.75
0.4589
94.37
0.4663
92.49
0.4112
T(T) = 0.3
99.66
0.6058
99.13
0.6088
97.64
0,5595
D(T) = 0.5
98.04
0.5377
95.36
0,5493
93.76
0.4928
PartialTSW (Ours)
((T) 0.9 P(T) 1.1
95.67
0,4515
94.39
0.4682
91.16
0.4024
92.03
0.4408
90.62
0.4533
89.34
0.4011
E.9.2 Additional Experimental Results
This section provides further details on model parameterization and presents complete set of results for all image-to-image translation directions.
Parameterization and Efficiency. As detailed in the implementation section, we adhered to the official configurations for the baseline methods. We utilized Gaussian Mixture Model (GMM) for ULightOT [28] and a flow-matching network for UOT-FM [22], ensuring a faithful and robust comparison.
For our method (PartialTSW) and the other standard OT baselines (SW [10], Db-TSW [80]), we employed the ResidualMLP architecture described in Appendix E.9. While GMM parameterization theoretically feasible for our method, we chose the neural network as it represents a more common, flexible, and standard approach for generative modeling tasks in recent literature.
This choice of parameterization not only standard but also highly efficient. As shown in Table our Resi didMMMMPPPPP approach is significantly more lightweight and faster per epoch than the complex models required by UOT-FM and ULightOT. This demonstrates that PartialTSW not only effective but also computationally efficient.
Results for All Translation Directions. In the main paper, our analysis centered on the Young-to- Adult (Y-A) translation task. This direction chosen as it represents the most significant class imbalance within the dataset. The class distribution of the pre- pre-processed FFHQ latent dataset is as follows: Young (approximately 15K Man, 23K Woman) and Adult (approximately 7K Man, 3.5K Woman).
To provide comprehensive analysis, we conducted additional experiments for all other possible translation directions: Woman-to-Man (W-M), Man-to-Wo to-onnn (M-W), and toilt-to-Yuugg (A-Y). The results, averaged over 5 independent runs, presented in Table 6.
50

---

These findings confirm that Partial-TSY consistently achieves superior trade-off between translation accuracy and perceptual similarity (LPIPS) across all translation settings. For instance, the W-M setting, Partial-TSW (with v(T) 0.9) achieves high accuracy of 95.67% while maintaining strong L.PISS of 0.4515. contrast, for UOT-FM to reach competitive accuracy (e.g., 81.50% with 0. 1),,, it incurs significantly worse LPIPS of 0.5198. This pattern, visible across the tasks, highlights our method's ability to find more effective and stable balance between the two competing objectives.
F Boarder Impacts
The introduction of the PartialTSW in this paper has a substantial societal impact by enhancing the precision and adaptability of optimal transport methods in various practical applications. This method can drive progress in numerous fields, such as healthcare, where better image processing techniques aid in more accurate medical imaging diagnostics, where enhanced generative models can lead to more sophisticated and creative outputs. Furthermore, the ability handle dynamic settings efficiently opens new possibilities for real-time data analysis in various sectors, including finance, logistics, and environmental monitoring.
the arts and entertainment industry,
and decision-maken Ultimately, the method contributes making advanced computational techniques more versatile and applicable to a broader range of real-works problems, thereby fostering innovation and improving societal well- beiigb
51

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed the limitations of the work in the Conclusion.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes].
52

---

We have provided full set of assumptions and complete proof for all theoretical results in Appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We have mentioned all information reproduce the main experimental results in Appendix E. All necessary details for reproducing the main experimental results are documented. This includes comprehensive descriptions network architectures, training procedures, and specific hyperparameters. Task-specific configurations are outlined in their respective subsections. The code is also included as supplementary material to further support reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibilitty In the case of closed-source models, it may be that access to the model is limited in
53

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We have provided data and code in the supplemental material, with detailed instructions reproduce the results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https: btttps://nipa cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: /nips. wwtppblicgguiieeeccocooooococcce http://guiiiseeeeeeeeee Ceeeee/CooSuubiieeee wwi.iissioopccyyc. for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We have provided all information about training and test details in Appendix E. We have included information on network architectures, optimizer types, batch sizes, number of training iterations, and the specific hyperparameters used in each experiment. In particular, Section E.7 describes the experimental setup for the point cloud experiments. Section E.8 focuses the robust generative model, and Section E.9 outlines the details for the image translation task.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Our experimental results are provided with error bars. We report statistical significance by including the mean and standard deviation over multiple runs in Table 2.
54

---

Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CL, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We have provided sufficient information about computing resources needed to reproduce the experiments in Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported didn't I make it into the paper).
paper (e.g., preliminary or failed experiments that NeurlPS Code of Ethics https://neur cc hccppbbbic/thhicssuuudeeeeeel
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the
Answer: [Yes]
Justification: The research conducted in the paper conforms, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
55

---

Answer:
[Yes]
Justification:
We have discussed Broader impacts Appendix.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets g.g code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All the assets used in the paper are properly cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
56

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi wwwdee com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The new assets provided in the paper are well documented, and the documen- tation is provided alongside the assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
57

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMS as any important, original, non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https / htmps://nnurrr //nnruurpp.... ecc/oonferecces/2225///////////e for what should or should not be described.
58


Reviews:
[{"review_id": "f68moYmeIa", "reviewer": "Reviewer_fEkm", "summary": "This paper extends the entropy partial transport framework on the tree sliced transport. Previously, the partial transport was introduced to address the issue of unbalanced mass, and the tree sliced transport was introduced to alleviate the computation burden of optimal transport and remedy the curse of dimensionality.  The proposed PartialTSW inherits the benefit of both frameworks. Theoretically the authors show that PartialTSW is a metric invariant to Euclidean transformations. In the experiments, the noise robustness of the proposed method is shown in the tasks of point cloud alignment, outlier rejection in generative modeling and image-to-image translation.", "questions": "1. Could the authors provide insights into the asymptotic behavior of PartialTSW with respect to the parameter $b$?\n3. Besides tuning $\\nu(\\mathcal{T})$ and fixing the slitting map to be softmax, how are the remaining regularization parameters chosen? \n5. Why are different baselines considered in different experiments?\n6. In the construction of tree systems, all lines pass through a common point $x$ in Algorithm 1. Is there a specific reason for this design choice? Additionally, how would the performance of the proposed method be affected if different $x_i$'s were randomly sampled instead?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The motivation is solid and clear. In practical applications, most distributions for matching are imbalanced, and efficiency is crucial. PartialTSW improves the robustness and efficiency of the classic optimal transport.\n2. The proposed PartialTSW is novel, and the theoretical justifications are well-written.\n3. The outlier robustness of PartialTSW is excellent and significant in experiments. I appreciate the experimental details given in Appendix E that help the reader understand the setting.\n\nWeaknesses:\nIncluding a discussion on sample complexity and projection complexity would enhance the completeness of the paper.", "was_revised": false, "final_justification": "", "rebuttal": "We address the concerns raised by the Reviewer in the **Weaknesses** and **Questions** sections as follows.\n\n---\n\n**W1. Including a discussion on sample complexity and projection complexity would enhance the completeness of the paper.**\n\n**Answer W1.** Similar to other applications of Monte Carlo (MC) estimation, the approximation error of the Tree-Sliced Wasserstein (TSW) distance decreases at a rate of $\\mathcal{O}(L^{-1/2})$, where $L$ denotes the number of independently sampled tree slices. This rate reflects the standard convergence behavior of MC-based approximations.\n\nConcretely, as discussed in the computational complexity analysis in Section 3 (lines 208), each tree slice in TSW contains $k$ projection lines. Therefore, with $L$ tree samples, the total number of projections is $Lk$. To ensure a fair comparison, we maintain the same total number of projections across all methods by adjusting $L$ accordingly, i.e., $L = \\text{(total number of projections)} / k$.\n\nIn the literature on TSW, it has been observed that using more than one line per tree slice ($k > 1$) enhances the capacity of the method to capture topological and structural features of the data. Thus, there is a natural trade-off: increasing $k$ improves topological expressiveness but requires reducing $L$ to preserve a fixed computational budget.\n\nIn our experiments, we tune the number of lines per tree slice $k$ based on empirical performance and adjust the number of tree slices $L$ accordingly to ensure that the total number of projections $Lk$ remains fixed. This approach enables us to explore the advantages of using richer tree structures through larger values of $k$ while maintaining a fair and consistent basis for comparison with other baseline methods under the same computational budget. Empirically, this strategy consistently leads to improved performance, highlighting the effectiveness of balancing structural expressiveness with projection efficiency.\n\n**Q1. Could the authors provide insights into the asymptotic behavior of PartialTSW with respect to the parameter $b$?**\n\n**Answer Q1.** In the Entropy Partial Transport (EPT) problem on tree metric spaces, the parameter $b$ appears as a weighting term in front of the transport cost:\n\n$W\\_m(\\mu, \\nu) = \\inf\\_{\\substack{\\gamma \\in \\Pi\\_{\\leq}(\\mu, \\nu)\\\\ \\gamma(\\mathcal{T} \\times \\mathcal{T}) = m}} \\left[ \\mathcal{F}\\_1(\\gamma\\_1 \\mid \\mu) + \\mathcal{F}\\_2(\\gamma\\_2 \\mid \\nu) + b \\int\\_{\\mathcal{T} \\times \\mathcal{T}} d\\_\\mathcal{T}(x, y) \\, \\gamma(dx, dy) \\right]$\n\nAs $b \\to \\infty$, the term $b \\int\\_{\\mathcal{T} \\times \\mathcal{T}} d\\_\\mathcal{T}(x, y) \\, \\gamma(dx, dy)$ dominates the objective. This corresponds to **penalizing the transport cost much more heavily** than the entropy terms $\\mathcal{F}\\_1$ and $\\mathcal{F}\\_2$. In this regime, the optimization will strongly prefer couplings that minimize the transport distance, possibly at the cost of higher entropy (i.e., less smooth or less spread-out plans).\n\nIn the limit $b \\to \\infty$, the optimizer would attempt to select a coupling $\\gamma$ that is highly concentrated around minimizing the distance $d_\\mathcal{T}(x, y)$, potentially leading to deterministic or near-deterministic couplings (subject to marginal constraints and the entropy penalty). \n\nAn analogous behavior is observed as $b \\to 0$. In conclusion,\n\n* **As $b \\to 0$**: the entropy terms dominate; the optimizer prefers smooth, spread-out couplings with less concern for transport distance.\n\n* **As $b \\to \\infty$**: the transport distance dominates; the optimizer prioritizes short transport paths, sacrificing entropy regularity.\n\nThus, $b$ controls the *bias-variance trade-off* between smoothness (entropy) and accuracy (transport cost). In practice, tuning $b$ balances regularization and fidelity.\n\n\n**Q2. Besides tuning $\\nu({\\mathcal{T}})$ and fixing the slitting map to be softmax, how are the remaining regularization parameters chosen?**\n\n**Answer Q2.** As discussed in Section 3.3, other regularization parameters (such as $a, b,$ and $\\lambda$) are not necessary for our method's primary application. We focus on computing gradients with respect to sample **locations**, while the sample **masses** are fixed and uniform (e.g., $1/n$) within a mini-batch.\n\nSince the model cannot learn to alter these masses, the parameters that would control mass-related penalties become irrelevant to the training process. This simplifies the formulation, leaving the transported mass ratio, $\\nu(\\mathcal{T})$, as the sole and intuitive hyperparameter for controlling the degree of partiality.\n\n**Q3. Why are different baselines considered in different experiments?**\n\n**Answer Q3.** The baselines were carefully selected for each experiment to ensure fair and meaningful comparisons based on both relevance and feasibility.\n\nThe gradient flow experiment (Sec 4.1) was designed as a focused comparison to isolate the benefit of our partial transport approach. We included only the direct balanced counterparts (SW and Db-TSW) to clearly demonstrate that the partial mechanism is a necessary advancement for handling outliers.\n\nThe robust generative modeling (Sec 4.2) and image-to-image translation (Sec 4.3) experiments provide a comprehensive evaluation against a wide range of state-of-the-art Unbalanced/Partial OT methods. We separated the baselines for two key technical reasons:\n* The methods from Sec 4.2 (e.g., USOT [1], SUOT [1], SPOT [2]) proved to be numerically unstable in the high-dimensional setting of Sec 4.3, frequently causing model divergence in our preliminary experiments.\n* Conversely, the methods from Sec 4.3 (UOT-FM and ULightOT) use highly specialized parameterizations (flow-matching networks and GMMs). Adapting them to the generative modeling task in Sec 4.2 would have required architectural changes, deviating from their original implementations and risking an unfair comparison.\n\nWe believe this approach provides the most rigorous evaluation, as PartialTSW is benchmarked against the most recent baselines within technical constraints of each task (e.g., PAWL [3] for Sec 4.2; UOT-FM [4] and ULightOT [5] for Sec 4.3).\n\n**Q4. In the construction of tree systems, all lines pass through a common point $x_0$ in Algorithm 1. Is there a specific reason for this design choice? Additionally, how would the performance of the proposed method be affected if different $x_0$'s were randomly sampled instead?**\n\n**Answer Q4.** The design of the tree system, in which each tree consists of $k$ lines emanating from a common root point, facilitates efficient algorithmic implementation. This design follows the approach in [6], where the authors provide a significantly more optimized implementation for such a tree structure, in contrast to the original formulation in [7] on Tree-Sliced Wasserstein (TSW).\n\nOur choice to adopt the tree-sliced framework, as introduced in [6, 7], is primarily driven by a foundational result of Caffarelli and McCann [8], which is summarized in lines 102–118 of our submission and elaborated in detail in Appendix B. Their result pertains to Entropic Partial Transport on tree-metric spaces. As outlined in Appendix B (lines 594–600), the construction involves the introduction of an auxiliary point $\\hat{s}$ into the space, which in turn necessitates working within a tree-metric structure rather than a conventional one-dimensional domain. This structural requirement naturally motivates the use of the tree-sliced framework over the traditional sliced approach.\n\nIt is important to note that in our implementation, the root point $x_0$ is not fixed. Rather, different $x_0$'s are sampled when computing the distance. As shown in Equation (17), the distance is estimated via Monte Carlo approximation by sampling $L$ tree systems, each consisting of $k$ lines. Each tree system is rooted independently, with its own randomly sampled $x_0$.\n\n---\n**References.**\n\n[1] Clément Bonet, et al., Slicing Unbalanced Optimal Transport. TMLR 2024.\n\n[2] Nicolas Bonneel and David Coeurjolly, SPOT: Sliced Partial Optimal Transport. ACM Transactions on Graphics (TOG) 2019.\n\n[3] Laetitia Chapel and Romain Tavenard, One for all and all for one: Efficient computation of partial Wasserstein distances on the line. ICLR 2025. \n\n[4] Luca Eyring et al., Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation. ICLR 2024.\n\n[5] Milena Gazdieva et al., Light Unbalanced Optimal Transport. NeuRIPS 2024.\n\n[6] Viet-Hoang Tran et al., Distance-based Tree-Sliced Wasserstein Distance. ICLR 2025.\n\n[7] Viet-Hoang Tran et al., Tree-Sliced Wasserstein Distance: A Geometric Perspective. ICML 2025.\n\n[8] Luis A Caffarelli and Robert J McCann, Free boundaries in optimal transport and Monge Ampere obstacle problems. Annals of Mathematics 2010.\n\n[9] Kilian Fatras et al., Unbalanced minibatch Optimal Transport; applications to Domain Adaptation. ICML 2021.\n\n[10] Chuyu Zhang et al., P$^2$OT: Progressive Partial Optimal Transport for Deep Imbalanced Clustering. ICLR 2024.\n\n[11] Henri De Plaen et al., Unbalanced Optimal Transport: A Unified Framework for Object Detection. CPVR 2023.\n\n---\n\nWe sincerely thank the Reviewer for the constructive feedback and for identifying the typographical errors in our submission. We will make the necessary revisions to address these issues. If the provided clarifications are deemed satisfactory, we would be grateful if the Reviewer could consider adjusting the evaluation accordingly. We remain at the Reviewer’s disposal and would be pleased to address any additional concerns in the subsequent stage of the discussion."}, {"review_id": "ZcAxiDoOIW", "reviewer": "Reviewer_nitj", "summary": "The paper considers approximations to optimal transport for measures on trees. More precisely, they consider a sliced approximation of the (possibly unbalanced) measures. Slicing is a popular way to reduce computational costs as computing optimal transport in 1D is often very easy. Partial transport relaxes the assumption that measures need to be probability measures and allows for unequal mass. This paper combines the two ideas to propose a tree-sliced partial optimal transport method.", "questions": "Questions:\n(1) What is *the* tree metric $d_{\\mathcal{T}}$? How does it relate to $\\mathrm{d}$?\n(2) I wasn't clear on what $\\alpha$ is. Unless I missed something, it first gets used just above eq (10) but isn't defined until after eq (13). And even then I wasn't sure what it means. Could you add more detail?\n(3) In Section 2.2, the paragraph on Tree System finishes with specifying the distribution $\\sigma$. It's just stated that this is a \"random sampling procedure\", is it uniform? Or any probability distribution over lines?\n(4) I didn't find the sentence \"we then compute the regularized transport cost $d_a(\\mu_{\\mathcal{T}},\\nu_{\\mathcal{T}})$ as in Equation (10)\" very clear. Eq (10) defines what I understand to be a 1-D projection. I feel there is a step missing in the logic. Could you add more detail.\n(5) Where does \"entropy\" in the name come from?\n\nMinor typo: Eq. (1) is the p-Wasserstein distance (not the 1-Wasserstein distance) as is written above.", "limitations": "There are a couple of sentences at the end of the paper.", "rating": 3, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Combining two ideas to produce new methods is popular in optimal transport where the geometry is often desirable but the computational costs and assumptions on the inputs (i.e. probability measures) can be prohibitive. So this is a nice idea that brings together two things, partial and slicing, to tree-based optimal transport distances.\n\nThe numerical results appear good. Although I wasn't sure why more comparisons weren't included in Section 4.1 (for example the methods in Section 4.2). It also wasn't clear how the tree graphs were constructed in the experimental section.", "was_revised": true, "final_justification": "This was the hardest of all the papers I had to review to make a judgement on. I suspect that from the slow responses of the other referees that their experience may be similar. And I understand how frustrating this is to the authors (and the AC). But it's a very long paper which I couldn't go through in detail in the time span available. I did go back through Appendix A and C which answered some of my questions. And it's perhaps reasonable that my other misunderstandings would be satisfied by going through everything else. I do accept the authors last response to my comments. However, without a better understanding I'm not comfortable increasing my score (and I'm not sure I think conferences are the best venue for long papers anyway), however I did decrease my confidence rating.", "rebuttal": "**General Response.** Based on the reviewer’s question, we believe there may be a lack of familiarity with the background of the Tree-Sliced Wasserstein (TSW) distance, as introduced in [1, 2]. In our submission, we provide a summary of the essential background in Section 2.2, with further technical details available in Appendices A and C. We offer here a brief introduction to the framework and encourage the reviewer to consult the aforementioned appendices for a deeper understanding.\n\nThe Tree-Sliced Wasserstein distance [2] was proposed as an extension of the traditional Sliced Wasserstein (SW) framework, where the projection domain, typically a line, is replaced by a more expressive structure known as a tree system. This enhancement has been empirically shown to capture richer structural information than a single line, while still admitting a closed-form approximation that allows for efficient computation.\n\nA tree system in $\\mathbb{R}^d$, roughly speaking, is a collection of multiple lines in $\\mathbb{R}^d$ that are connected through a tree structure. This configuration induces a tree-metric space, on which the optimal transport problem admits closed-form solutions, as detailed in Section 2.1 and Appendix A.\n\nWithin this tree-metric space, the TSW framework differs fundamentally from the SW approach. While SW projects the entire mass at each support point of a distribution onto a single line, TSW distributes the mass at each support point into $k$ parts, each assigned to one of the $k$ lines in the tree system. This process takes into account the geometric relationship between the support points and the structure of the tree system, leading to a significantly different projection mechanism.\n\nOnce projected, the TSW framework leverages the closed-form solution for optimal transport on tree-metric spaces to compute the TSW distance efficiently. This makes the computational complexity of TSW comparable to that of SW. However, TSW has been shown to yield superior empirical performance, making it a promising alternative to the traditional SW distance and motivating further study of the tree-sliced approach.\n\n---\n\nWe address the concerns raised by the Reviewer in the **Weaknesses** and **Questions** sections as follows.\n\n---\n\n**Q1. What is the tree metric $d_\\mathcal{T}$? How does it relate to $d$?**\n\n**Answer Q1.** Here, the $d_\\mathcal{T}$ is the tree metric of the metric space $\\mathcal{T}$. We are not sure what $d$ that the Reviewer mentioned, and assume it is the $d_a$ in the Equation (9). Heree, $d_a$ is a metric on the space of measures, that is studied in literature.\n\n**Q2. I wasn't clear on what $\\alpha$ is. Unless I missed something, it first gets used just above eq (10) but isn't defined until after eq (13). And even then I wasn't sure what it means. Could you add more detail?**\n\n**Answer Q2.** We acknowledge that, during the compression of the paper to meet the required length constraints, we inadvertently omitted the definition of the notion $\\alpha$. The map $\\alpha$, referred to as the splitting map, is an element of the space $\\mathcal{C}(\\mathbb{R}^d \\times \\mathbb{T}^d_k, \\Delta_{k-1})$. Its definition is provided in Appendix C, specifically in Appendix C.2 (line 685).\n\nThis map encodes the mass-splitting mechanism: for each support point $x \\in \\mathbb{R}^d$, the density is distributed across the $k$ lines of a tree system $\\mathcal{T} \\in \\mathbb{T}^d_k$ according to the probability vector $\\alpha(x, \\mathcal{T}) \\in \\Delta_{k-1}$. While the framework allows for any choice of $\\alpha$ in $\\mathcal{C}(\\mathbb{R}^d \\times \\mathbb{T}^d_k, \\Delta_{k-1})$, the choice proposed in [1], based on the relative distances between the point $x$ and the lines of the tree system, has been shown to perform well in practice. This motivates the specific form of $\\alpha$ used in our submission, as defined in Equation (13).\n\nWe thank the reviewer for identifying this omission and will revise the submission accordingly to clarify the definition and role of $\\alpha$.\n\n**W2. It also wasn't clear how the tree graphs were constructed in the experimental section.**\n\n**Q3. In Section 2.2, the paragraph on Tree System finishes with specifying the distribution $\\sigma$. It's just stated that this is a \"random sampling procedure\", is it uniform? Or any probability distribution over lines?**\n\n**Answer W2+Q3.** The detailed construction of the distribution $\\sigma$ is provided in Appendix C.1. This distribution is defined based on the procedure used to construct the tree system, specifically through the sampling of the root point and the lines that form the tree structure.\n\n**Q4. I didn't find the sentence \"we then compute the regularized transport cost $\\mathcal{T}_{\\epsilon}(\\mu, \\nu)$ as in Equation (10)\" very clear. Eq (10) defines what I understand to be a 1-D projection. I feel there is a step missing in the logic. Could you add more detail.**\n\n**Answer Q4.** We acknowledge the typo in the original sentence and thank the reviewer for pointing it out. The corrected sentence should read: “We then compute the regularized transport cost $\\mathcal{T}\\_{a}(\\mu\\_\\mathcal{T}, \\nu\\_\\mathcal{T})$ as in Equation (9).” Here, Equation (9) explicitly defines the computation of $\\mathcal{T}\\_{a}(\\mu\\_\\mathcal{T}, \\nu\\_\\mathcal{T})$. We kindly invite the reviewer to refer to this equation for clarification.\n\n**Q5. Where does \"entropy\" in the name come from?**\n\n**Answer Q5.** The term “Entropy” originates from Equations (5) and (6), where the Partial Transport problem is formulated with regularization terms expressed as weighted relative entropies.\n\n**W1. Although I wasn't sure why more comparisons weren't included in Section 4.1 (for example the methods in Section 4.2)**\n\n**Answer W1.** The baselines were carefully selected for each experiment to ensure fair and meaningful comparisons based on both relevance and feasibility.\n\nThe gradient flow experiment (Sec 4.1) was designed as a focused comparison to isolate the benefit of our partial transport approach. We included only the direct balanced counterparts (SW and Db-TSW) to clearly demonstrate that the partial mechanism is a necessary advancement for handling outliers.\n\nThe robust generative modeling (Sec 4.2) and image-to-image translation (Sec 4.3) experiments provide a comprehensive evaluation against a wide range of state-of-the-art Unbalanced/Partial OT methods. We separated the baselines for two key technical reasons:\n* The methods from Sec 4.2 (e.g., USOT, SUOT, SPOT) proved to be numerically unstable in the high-dimensional setting of Sec 4.3, frequently causing model divergence in our preliminary experiments.\n* Conversely, the methods from Sec 4.3 (UOT-FM and ULightOT) use highly specialized parameterizations (flow-matching networks and GMMs). Adapting them to the generative modeling task in Sec 4.2 would have required architectural changes, deviating from their original implementations and risking an unfair comparison.\n\nWe believe this approach provides the most rigorous evaluation, as PartialTSW is benchmarked against the most recent baselines within technical constraints of each task (e.g., PAWL [3] for Sec 4.2; UOT-FM [4] and ULightOT [5] for Sec 4.3).\n\n---\n\n**References.**\n\n[1] Viet-Hoang Tran et al., Distance-based Tree-Sliced Wasserstein Distance. ICLR 2025.\n\n[2] Viet-Hoang Tran et al., Tree-Sliced Wasserstein Distance: A Geometric Perspective. ICML 2025.\n\n[3] Laetitia Chapel and Romain Tavenard, One for all and all for one: Efficient computation of partial Wasserstein distances on the line. ICLR 2025.\n\n[4] Luca Eyring et al., Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation. ICLR 2024.\n\n[5] Milena Gazdieva et al., Light Unbalanced Optimal Transport. NeurIPS 2024.\n\n---\n\nWe sincerely thank the reviewer for the constructive feedback and for identifying the typographical errors in our submission. We will make the necessary revisions to address these issues. If the provided clarifications are deemed satisfactory, we would be grateful if the reviewer could consider adjusting the evaluation accordingly. We remain at the reviewer’s disposal and would be pleased to address any additional concerns in the subsequent stage of the discussion."}, {"review_id": "4EgCsgO9FJ", "reviewer": "Reviewer_uZKD", "summary": "Tree-based approach to compute Sliced Partial Optimal Transport.", "questions": "1) How is the parametrization of the proposed approach tool different from the comparison methods, such as ULIghtOT, which uses a mixture of Gaussians, not neural networks? This makes the comparison unfair. Can you parameterize your method using a mixture of Gaussians?\n\n2) How do the results depend on the parameter k? Can we control this parameter while *building a tree*?\n\n3) Does PartialTSW have lower transport plan costs or Wasserstein distances between the target and generated distributions than other unbalanced methods?\n\n4) Why did the authors only consider young-to-adult translation for the image generation task? The same dataset could be used for adult-to-young, man-to-woman, and woman-to-man experiments without changing the setup.\n\n5) More diverse datasets should be included. In which real-world scenarios could your approach be used?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths:** I would highlight the clarity of the paper as one of its strengths. It is well-written and logically structured, making it easy to follow. The approach of efficiently solving partial sliced optimal transport using tree-based priors is also novel and promising for easily computing unbalanced OT.  I appreciate the authors' experiments on the dependence of unbalanced/tau parameters, which are significant for unbalanced/partial settings. The proposed method clearly demonstrates stability and dependence on the parameters, and, in comparison to other OT approaches, it provides an interpretable value between 0 and 1.\n\n**Weaknesses:** The only change from [73, 77] is the switch to an Unbalanced OT objective, which minimizes the methodological contribution. Aside from the theorem that PartialTSW is an E(d)-invariant metric on M(Rd), the theoretical contributions are limited. Besides the theoretical contribution, the practical justification is also limited. The authors provided several experiments. I think the approach needs a more expressive, high-dimensional justification of the proposed method. The robust generative modeling experiments are interesting, but they are very artificial (e.g., MNIST). and do not cover non-real-world scenarios This makes the justification for the proposed method unclear.", "was_revised": true, "final_justification": "My concerns are satisfied, rising my score.", "rebuttal": "We address the concerns raised by the Reviewer in the **Weaknesses** and **Questions** sections as follows.\n\n---\n\n**W1. The only change ... minimizes the methodological contribution.**\n\n**W2. Aside from the theorem ... theoretical contributions are limited.**\n\n**Answer W1+W2.** We would like to clarify that we view our extension of TSW to the partial setting as a substantive contribution, for the following reasons:\n\n- The Entropy Partial Transport (EPT) problem on tree metric spaces has been carefully studied, primarily for its computational advantages in large-scale settings. However, prior applications have been limited to distributions supported on graph- or tree-structured spaces, which significantly restricts their applicability to broader settings involving Euclidean measures or generative tasks. Extending EPT to Euclidean measures through the Tree-Sliced Wasserstein (TSW) framework requires a new conceptual bridge—insight that, to the best of our knowledge, has not been previously established.\n\n- Our work provides both theoretical and experimental contributions that detail the adoption of the EPT formulation within the TSW framework for Euclidean data. This includes the derivation of a valid metric between measures, $\\text{E}(d)$-invariance, supported by closed-form approximations and efficient implementations. All theoretical claims are accompanied by rigorous proofs. \n\n- While prior works have applied unbalanced or partial transport to sliced variants of Wasserstein distances (see [1], [2]), the EPT problem on tree metrics introduces unique structural requirements—specifically, the need to augment the support with an additional point (as noted in the Reviewer’s question). This structural modification cannot be accommodated within the standard sliced setting and instead necessitates the full TSW framework. We address this by proposing a closed-form approximation of our distance (see Equation (9)), which leads to significant computational advantages. Empirically, our method outperforms prior approaches in terms of runtime.\n\nFor these reasons, we believe our contribution is substantive and non-incremental, offering both theoretical novelty and practical utility within the broader context of optimal transport on Euclidean domains.\n\n**W3. Besides the theoretical contribution, the practical justification ...justification for the proposed method unclear.**\n\n**Q5. More diverse datasets ...real-world scenarios could your approach be used?**\n\n**Answer W3+Q5.** We appreciate the reviewer's positive feedback on our generative modeling experiment. We designed this experiment on MNIST as a clear, illustrative study to visually demonstrate the qualitative advantages of our method's robustness. While this setup is simplified for clarity, the principle of using PartialTSW to build resilience against outlier data is a general one, applicable to more complex datasets and other modalities.\n\nFor a high-dimensional justification, we presented an image-to-image translation task in Section 4.3 ($d=512$), adapted from ULightOT [5]. The strong performance of PartialTSW in this experiment, particularly its ability to handle significant class imbalance, directly demonstrates its effectiveness and scalability in a challenging, non-trivial setting.\n\nThe applicability of PartialTSW extends beyond our presented experiments. As a robust and differentiable loss function, it is well-suited for tasks like domain adaptation [9] and self-supervised learning [10]. We note, however, that it is not designed for assignment-based problems (e.g., object detection [11]), as it does not compute an explicit transport plan.\n\n**Q1. How is the parametrization ... using a mixture of Gaussians?**\n\nWe thank the reviewer for this important question regarding the parameterization in the image-to-image translation experiments.\n\nFor the baselines, we faithfully followed the official settings from their respective papers, using a Gaussian Mixture Model for ULightOT [5] and a flow-matching network for UOT-FM [4], to ensure a strong and fair representation of their methods. These details are presented in Appendix E.8.\n\nFor our method (PartialTSW) and the other OT-based baselines (SW, Db-TSW [6]), we used a standard neural network. While parameterizing our method with a GMM is theoretically possible, we opted for a neural network as it is a more common and flexible choice in the current literature for generative tasks.\n\nWe respectfully argue that this comparison is not unfair to the baselines; in fact, our chosen parameterization is significantly more efficient. The following table summarizes the parameter counts and per-epoch training times:\n\n|Method|Number of Parameters|Time (s/epoch)|\n|:-|:-|:-|\n|SW|526,337|6|\n|Db-TSW|526,337|25|\n|**PartialTSW (Ours)**|**526,337**|**25**|\n|UOT-FM|788,224|35|\n|ULightOT|5,263,380|60|\n\nAs shown, our approach achieves its results using fewer parameters and faster training time per epoch than ULightOT and UOT-FM. This demonstrates that our method is not only effective but also computationally efficient.\n\n**Q2.  How do the results ... control this parameter while building a tree?**\n\nThe number of lines per tree, $k$, is a hyperparameter that is set during the tree construction process, which is detailed in Appendix C.1.\n\nRegarding its effect on performance, our empirical analysis shows that the best results are achieved by prioritizing a large number of trees ($L$) over a high tree complexity ($k$). Specifically, a low value for $k$ (e.g., $k \\in [2, 10]$) provides the best trade-off between the stability of the distance estimate and overall computational speed. We provide a more detailed treatment of the interplay between $L$ and $k$ in our response to Q4 from Reviewer d8LY.\n\n**Q3. Does PartialTSW have lower transport plan costs or Wasserstein distances between the target and generated distributions than other unbalanced methods?**\n\n**Answer Q3.** We are not entirely certain about the motivation behind the question raised by the Reviewer. In our setting, PartialTSW—like other baselines—is used as a loss function to be minimized, not for direct magnitude comparison. Final evaluations are typically based on a universal metric, ideally from unbalanced OT.\n\nTherefore, comparing the absolute values of distances computed by PartialTSW and other baselines is not necessarily meaningful. In designing a distance (or discrepancy) between probability measures, the primary goal is to use it as a loss function in applications such as generative modeling or distribution alignment, where the objective is to minimize divergence between the model and empirical distributions.\n\nThis objective is well-defined when the discrepancy is a true metric—i.e., it satisfies the identity of indiscernibles: $\\text{discrepancy}(\\mu,\\nu) = 0$ if and only if $\\mu =\\nu$. Without this property, minimizing the loss may not guarantee distributional equivalence, making the optimization ambiguous. Hence, this property is essential for both theoretical validity and interpretability.\n\nEach proposed distance—including the baselines and our PartialTSW—induces a different loss landscape, varying in smoothness, optimization difficulty, and geometric alignment.  Empirically, PartialTSW yields a favorable loss landscape that leads to improved convergence and performance across several tasks.\n\n**Q4. Why did the authors ... without changing the setup.**\n\n**Answer Q4.** For the main paper, we focused on the Young-to-Adult (Y→A) translation task as it presents the most significant class imbalance in the dataset. The class distribution is as follows:\n\n||Man|Woman|\n|:-|:-:|:-:|\n|**Young**|15K|23K|\n|**Adult**|7K|3.5K|\n\nAs suggested by the Reviewer, we conducted additional experiments for all translation directions. The results, averaged over 5 runs, are summarized in the table below. They confirm that **Partial-TSW consistently achieves a better trade-off between task accuracy and perceptual similarity (LPIPS)** across all settings.\n\nFor instance, in the Woman-to-Man (W→M) setting, Partial-TSW with $\\nu(\\mathcal{T})=0.9$ achieves a high accuracy of 95.67% while maintaining an excellent LPIPS score of 0.4515. In contrast, for UOT-FM to achieve a competitive accuracy, it requires a large regularization ($\\epsilon=0.1$), which results in a significantly worse LPIPS of 0.5198. This highlights our method's ability to find a more effective balance.\n\n|Method||W$\\rightarrow$M Acc (%) $\\uparrow$|W$\\rightarrow$M LPIPS $\\downarrow$|M$\\rightarrow$W Acc (%) $\\uparrow$|M$\\rightarrow$W LPIPS $\\downarrow$|A$\\rightarrow$Y Acc (%) $\\uparrow$|A$\\rightarrow$Y LPIPS $\\downarrow$|\n|:-|:-|:-:|:-:|:-:|:-:|:-:|:-:|\n|**SW**||93.95|0.4418|91.68|0.4546|89.82|0.4041|\n|**Db-TSW**||94.13|0.4436|92.07|0.4546|89.58|0.4022|\n|**UOT-FM**||||||||\n||$\\epsilon=0.0005$|49.21|0.3914|85.52|0.4269|79.77|0.3836|\n||$\\epsilon=0.005$|70.69|0.4531|94.19|0.4749|92.91|0.4129|\n||$\\epsilon=0.05$|80.28|0.4899|95.40|0.5106|97.40|0.4693|\n||$\\epsilon=0.1$|81.50|0.5198|97.91|0.5369|98.43|0.4828|\n|**ULightOT**||||||||\n||$\\tau=50.0$|76.36|0.4102|85.91|0.4086|84.31|0.3452|\n||$\\tau=250.0$|86.36|0.4466|92.81|0.4516|90.49|0.3906|\n||$\\tau=1000.0$|88.07|0.4557|93.91|0.4626|92.00|0.4060|\n||$\\tau=10000.0$|88.75|0.4589|94.37|0.4663|92.49|0.4112|\n|**PartialTSW (Ours)**||||||||\n||$\\nu(\\mathcal{T})=0.3$|99.66|0.6058|99.13|0.6088|97.64|0.5595|\n||$\\nu(\\mathcal{T})=0.5$|98.04|0.5377|95.36|0.5493|93.76|0.4928|\n||$\\nu(\\mathcal{T})=0.9$|95.67|0.4515|94.39|0.4682|91.16|0.4024|\n||$\\nu(\\mathcal{T})=1.1$|92.03|0.4408|90.62|0.4533|89.34|0.4011|\n\n---\n\n**References.** Due to space constraints, we refer the Reviewer to the references provided at the end of our response to **Reviewer fEkm**.\n\n---\n\nWe thank the Reviewer for the constructive feedback and will revise the manuscript accordingly. If our clarifications are satisfactory, we kindly hope the evaluation may be reconsidered. We remain available to address any further concerns in the next stage of the discussion."}, {"review_id": "zdbHStFsOn", "reviewer": "Reviewer_d8LY", "summary": "The paper introduces Tree-Sliced Entropy Partial Transport (PartialTSW), obtained by averaging an entropy–regularised partial transport cost over sets of random tree metrics. It propose to compute the partial transport problem on a tree as a balanced OT problem between modified measures. The distance is shown to satisfy the metric proprieties and inherits rotation-invariance from the random‐tree construction. Experiments includes synthetic point-cloud interpolation with noise, image latent-space auto-encoding and class-imbalanced image translation, reporting speed-ups over Sinkhorn-based unbalanced OT and robustness to added outliers.", "questions": "Without a quantitative cost–accuracy curve the runtime tables (Fig. 2, Table 1) risk comparing methods at different error tolerances. Could \n you please either provide (i) a variance/concentration bound for the Monte-Carlo estimator and a bias analysis for finite 𝑘, or (ii) empirical convergence plots so readers can see how many projections are needed for stable distances. \n\nCould you include the partial-transport baselines (SUOT, SPOT, USOT, PAWL) in the point-cloud and generative experiments, with the same transported-mass budget ν(T)?\n\nIs there any possible to add an outlier-robustness theorem or, at minimum, systematic contamination sweeps that hold ν(T) fixed across methods? Without such evidence, the robustness claim is week.\n\nCould you supply guidance for choosing L and k? An empirical plot of the estimation error against Lk would enable users to trade accuracy for speed.", "limitations": "na", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Quality: The algebraic reduction from partial to balanced OT on a tree and proofs of metric properties looks good to me. Implement use cases are enough. But the numerical accuracy is not controlled across baselines; runtime comparisons therefore lack a common error tolerance. No variance or convergence analysis is provided for the Monte-Carlo estimator.\n\nClarity: The manuscript is generally well written. But the role and tuning of the transported-mass parameter ν(T) are only discussed heuristically.\n\nSignificance: The O(nlogn) time complexity is appealing by matching the complexity of balanced TSW, and runs faster and appears more robust than the rest of baselines in experiments. However, without cost–accuracy curves or formal robustness lemma the empirical impact is uncertain.\n\nOriginality: Extending TSW to the partial setting is incremental. I think the novel step is the balanced-reduction identity on trees; all other components (random-tree averaging, robustness via partial mass) follow from prior work on TSW and partial OT.\n\nOverall, the work lacks both a theoretical error analysis and decisive empirical evidence that PartialTSW is more robust or faster at equal accuracy than existing partial or unbalanced OT solvers. I therefore lean to reject.", "was_revised": true, "final_justification": "I have reviewed the response and believed they addressed my concerns. Raised my score to 4.", "rebuttal": "We address the concerns raised by the Reviewer in the **Weaknesses** and **Questions** sections as follows.\n\n---\n\n**W1. The role and tuning ... discussed heuristically.**\n\n**Answer W1**. We argue that tuning the transported-mass parameter $\\nu(\\mathcal{T})$ is a key practical advantage of our method due to its interpretability and principled range.\n\nUnlike baseline methods where partiality is controlled by an abstract regularization term, $\\nu(\\mathcal{T})$ directly corresponds to the **fraction of mass to be transported**. This allows for an intuitive, data-driven tuning process. For example, in a setting with an estimated 20% outlier or noise ratio, a principled starting point for tuning is $\\nu(\\mathcal{T}) \\approx 0.8$.\n\nFurthermore, this provides a constrained and interpretable search space for the hyperparameter, typically within $(0, 1]$. This contrasts sharply with methods like ULightOT [1] or UOT-FM [2], where the equivalent parameters ($\\tau$ or $\\epsilon$) are not directly interpretable and must be tuned across several orders of magnitude to achieve meaningful differences in partiality, as shown in our experiments (Table 2).\n\n**Q1. Without a quantitative cost–accuracy ... how many projections are needed for stable distances.**\n\n**W2. The $O(n \\log n)$ time complexity ... empirical impact is uncertain.**\n\n**Answer Q1+W2**. We thank the reviewer for raising this crucial point about the accuracy-cost trade-off. We agree that runtime comparisons are only meaningful at a similar error tolerance.\n\nTo address this, we first clarify the roles of our method's parameters. For a number of lines ($k$), a ground-truth distance $PartialTSW(\\mu, \\nu)$ is defined. Our method then uses a **Monte Carlo estimation over the number of random trees ($L$)** to approximate this distance. The reviewer's concern about the stability of this estimator is well-founded.\n\nAs requested, we provide an empirical analysis of this estimator's convergence. The table below shows the **Coefficient of Variation (CoV = $\\sigma/\\mu$)**, a normalized measure of stability, as a function of both $L$ (the Monte Carlo parameter) and $k$ (a structural parameter).\n\n|$L$|$k=5$|$k=10$|$k=100$|\n|:-|:-:|:-:|:-:|\n|10|0.1098|0.1016|0.0456|\n|50 | 0.0526 | 0.0252|0.0183|\n|100|0.0263|0.0239|0.0091|\n|500|0.0105 |0.0133|0.0070|\n|1000|0.0076|0.0070|0.0041|\n|10000|0.0029| 0.0018|0.0023|\n|20000|0.0016|0.0024|0.0010|\n\nThe results empirically verify that for any fixed $k$, the estimator's stability improves as $L$ increases, aligning with the theoretical $O(L^{-1/2})$ Monte Carlo convergence rate. For instance, at $k=5$, increasing $L$ from 10 to 1000 reduces the CoV by over 14x. In our large-scale experiments, we use a high number of projections ($L \\ge 1000$) to ensure a stable estimate. We will add this analysis to the appendix to guide practitioners.\n\nFor the runtime comparisons in Figure 2 and Table 1, we ensured a fair comparison across all sliced-based methods. We used an identical number of projections ($L$) for all methods. Additionally, for the tree-based methods (TSW and our PartialTSW), the tree complexity parameter ($k$) was held constant.\n\nMore importantly, your question highlights a key advantage of our method. The computational complexity of PartialTSW is **independent** of its robustness parameter, $\\nu(\\mathcal{T})$. In contrast, the complexity of baselines like PAWL [3] is directly tied to their robustness parameter (the number of samples $k$ to transport). This means that for such baselines, increasing robustness inherently increases computational cost. PartialTSW does not face this trade-off. The fact that it is both highly robust and empirically among the fastest methods underscores its practical significance.\n\n**W3. Extending TSW to the partial setting is incremental. I think the novel step is the balanced-reduction identity on trees; all other components (random-tree averaging, robustness via partial mass) follow from prior work on TSW and partial OT.**\n\n**Answer W3**. We would like to clarify that we view our extension of TSW to the partial setting as a substantive contribution, for the following reasons:\n\n- The Entropy Partial Transport (EPT) problem on tree metric spaces has been carefully studied, primarily for its computational advantages in large-scale settings. However, prior applications have been limited to distributions supported on graph- or tree-structured spaces, which significantly restricts their applicability to broader settings involving Euclidean measures or generative tasks. Extending EPT to Euclidean measures through the TSW framework requires a new conceptual bridge, insight that, to the best of our knowledge, has not been previously established.\n\n- Our work provides both theoretical and experimental contributions that detail the adoption of the EPT formulation within the TSW framework for Euclidean data. This includes the derivation of a valid metric between measures, $\\text{E}(d)$-invariance, supported by closed-form approximations and efficient implementations. All theoretical claims are accompanied by rigorous proofs.\n\n- While prior works have applied unbalanced or partial transport to sliced variants of Wasserstein distances (see [4], [5]), the EPT problem on tree metrics introduces unique structural requirements, specifically, the need to augment the support with an additional point (as noted in the Reviewer’s question). This structural modification cannot be accommodated within the standard sliced setting and instead necessitates the full TSW framework. We address this by proposing a closed-form approximation of our distance (see Equation (9)), which leads to significant computational advantages. Empirically, our method outperforms prior approaches in terms of runtime.\n\nFor these reasons, we believe our contribution is substantive and non-incremental, offering both theoretical novelty and practical utility within the broader context of optimal transport on Euclidean domains.\n\n**Q2. Could you include ... same transported-mass budget $\\nu({\\mathcal{T}})$?**\n\n**Q3. Is there any possibility ... robustness claim is weak.**\n\n**Answer Q2+Q3.** Regarding **Q2**, we agree that comparing methods with a matched transported-mass budget is ideal. However, this is often not feasible due to fundamental differences in how the baselines are formulated. While our PartialTSW allows for a flexible, continuous mass budget via the $\\nu(\\mathcal{T})$ parameter, competing methods control partiality indirectly:\n* Methods like **SPOT** [4] and **PAWL** [3] rely on selecting an integer number of samples, $k$.\n* For **USOT** [5] and **SUOT** [5], while they can theoretically be re-weighted, we found their available implementations did not support this direct control in our experiments, necessitating that partiality be managed via their regularization parameters ($\\rho_1, \\rho_2$).\n\nBecause a direct matching of the mass budget is not possible, we focused on finding the optimal robustness setting for each method.\n\nThis directly addresses **Q3**. While we have not included a formal outlier-robustness theorem, we did perform a **systematic hyperparameter sweep for all baselines** (detailed in Appendix E.4 and E.7.2) to find the settings that best reject outliers.\n\nOur findings show that even at their optimal settings, competing methods struggle. Baselines like USOT and SUOT, despite extensive tuning of their regularization parameters, could not completely ignore outliers, demonstrating the limitations of their indirect control mechanism. While a recent method like PAWL can achieve strong outlier rejection, it is slower than our method. In contrast, our experiments demonstrate that PartialTSW provides a direct, efficient, and effective mechanism for handling outliers.\n\n**Q4. Could you supply guidance for choosing $L$ and $k$? ... trade accuracy for speed.**\n\n**Answer Q4.** As established in our answer to **Q1+W2**, the Monte Carlo estimation error is a function of the number of trees, $L$, while the computational complexity scales linearly with both $L$ and the number of lines per tree, $k$.\n\nOur key recommendation for a fixed computational budget $O(Lk)$ is to prioritize a **large number of trees (high $L$, low $k$)**. The reasoning is that the primary source of error is the Monte Carlo variance, which is most effectively reduced by increasing $L$. Our empirical analysis in the Coefficient of Variation table (in our response to **Q1**) supports this, showing that increasing $L$ consistently provides the largest and most stable reduction in estimation error.\n\nBased on our experiments, we offer the following practical guidance: for large-scale tasks like the ones presented in our paper, we found a low number of lines **$k \\in [2, 10]$** combined with a large number of trees, **$L \\in [1000, 10000]$**, provides the best trade-off between speed and stable performance. We will add this explicit guidance to the appendix.\n\n---\n\n**References.**\n\n[1] Milena Gazdieva et al., Light Unbalanced Optimal Transport. NeurIPS 2024.\n\n[2] Luca Eyring et al., Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation. ICLR 2024.\n\n[3] Laetitia Chapel and Romain Tavenard, One for all and all for one: Efficient computation of partial Wasserstein distances on the line. ICLR 2025. \n\n[4] Nicolas Bonneel and David Coeurjolly, SPOT: Sliced Partial Optimal Transport. ACM Transactions on Graphics (TOG) 2019.\n\n[5] Clément Bonet, et al., Slicing Unbalanced Optimal Transport. TMLR 2024.\n\n---\n\nWe sincerely thank the reviewer for the constructive feedback and for identifying the typographical errors in our submission. We will make the necessary revisions to address these issues. If the provided clarifications are deemed satisfactory, we would be grateful if the reviewer could consider adjusting the evaluation accordingly. We remain at the reviewer’s disposal and would be pleased to address any additional concerns in the subsequent stage of the discussion."}]

Author rebuttal:

Meta-review: This paper proposes a new method for partial optimal transport (with entropic divergence) using tree-sliced
OT. The authors prove metric properties of the method (with  E(d)-invariance)
and provide a closed form approximation for computing it. They evaluate the
method on point cloud gradient flow, robust generative modeling and unbalanced image to
image translation and their proposed distance is competitive with the selected
baselines.

The reviewers found the paper interesting and rigorous but had some
concerns about the originality of the method (combination of Tree Sliced and
Unbalanced OT), the clarity of the presentation
(lacking some details), and the limited numerical experiments (missing
sensitivity to L and k independently, missing more image translation exp.).

The authors did a nice and detailed response to all those points with clarifications
and additional experiments. The reviewers appreciated those clarifications and
new results (L/K, image translation) and now agree that the paper can be accepted
if the authors include all those in the final version.

I agree with the reviewers that this paper is interesting and has a good
theoretical contribution and provides a new method for partial tree-sliced
optimal transport. I recommend an accept and expect the final version to include
the discussed improvements.
Decision: Accept (poster)