Tree-Sliced Wasserstein Distance with Nonlinear Projection
Authors: Thanh Tran
Hoang V. Tran
Thanh Chu
Huyen Trang Pham
Laurent Ghaoui
Tam Le
Tan Minh Nguyen
Venue: ICML 2025 poster (2025)
Keywords: tree-sliced wasserstein distance
optimal transport
Primary area: deep_learning
TL;DR: We propose a novel nonlinear projectional framework for the Tree-Sliced Wasserstein distance, substituting the linear projections in earlier versions with general projections

Abstract: Tree-Sliced methods have recently emerged as an alternative to the traditional Sliced Wasserstein (SW) distance, replacing one-dimensional lines with tree-based metric spaces and incorporating a splitting mechanism for projecting measures. This approach enhances the ability to capture the topological structures of integration domains in Sliced Optimal Transport while maintaining low computational costs. Building on this foundation, we propose a novel nonlinear projectional framework for the Tree-Sliced Wasserstein (TSW) distance, substituting the linear projections in earlier versions with general projections, while ensuring the injectivity of the associated Radon Transform and preserving the well-definedness of the resulting metric. By designing appropriate projections, we construct efficient metrics for measures on both Euclidean spaces and spheres. Finally, we validate our proposed metric through extensive numerical experiments for Euclidean and spherical datasets. Applications include gradient flows, self-supervised learning, and generative models, where our methods demonstrate significant improvements over recent SW and TSW variants.

Tree-Sliced Wasserstein Distance with Nonlinear Projection
Thanh Tran
Viet-Hoang Tran
Thanh Chu Trang Pham Laurent El Ghaoui Tam Le 14
Tan M. Nguyen f2
Abstract
learning (Nguyen al., 2021b; Bunne et al., 2022; Hua et al., 2023; Fan et al., 2022; Le et al., 2024a;b; Kessler et al., 2025; Chapel et al., 2025; Chapel & Tavenard, 2025), multimodal data analysis (Park et al., 2024; Luong et al., 2024), computer vision and graphics (Lavenant et al., 2018;
Tree-Sliced methods have recently emerged an alternative to the traditional Sliced Wasserstein (SW) distance, replacing one-dimensional lines with tree-based metric spaces and incorporating a splitting mechanism for projecting measures. This approach enhances the ability to capture the topological structures of integration domains in Sliced Optimal Transport while maintaining low computational costs. Building this foundation, we propose a novel nonlinear projectional frame- work for the Tree-Sliced Wasserstein (TSW) dis- tance, substituting the linear projections in earlier versions with general projections, while ensur- ing the injectivity of the associated Radon Trans- form and preserving the well-definedness of the resulting metric. By designing appropriate pro- jections, we construct efficient metrics for mea- sures on both Euclidean spaces and spheres. Fi- nally, we validate our proposed metric through extensive numerical experiments for Euclidean and spherical datasets. Applications include gra- dient flows, sulf-super learning, and gener- ative models, where our methods demonstrate significant improvements over recent SW and TSW variants. The code is publicly available at /atpps://githhe www/ htmm/hhanngtt200// Non1 inea rTWW.
Nguyen et al., 202 a; Saleh et al., 2022: Rabin et al., 2011; Solomon et al., 2015; Vu et al.. 2025), statistics (Mena &
2019; Weed & Berthet, 2019; Wang et al.,
2022; Pham et al., 2024; Liu et al., 2022; Nguyen et al., 2022; Nietert et al., 2022). By utilizing the closed-form solution of one-dimensional OT problems, SW substantially reduces the computational complexity typically associated with OT (Rabin et al., 2011; Bonneel et al., 2015; Peyré et al., 2019).
Related work. Several enhancements have been proposed to improve different aspects of the SW framework, including optimizing the sampling process (Nguyen et al., 2020; Nad- jahi et al., 2021; Nguyen et al., 2024a), selecting optimal projection directions (Deshpande et al., 2019), and refining the projection mechanism (Kolouri et al., 2019; Chen et al., 2022; Bonet et al., 2023).
The Tree-Sliced framework (Tran al., 2025c;a;b) has re- cently emerged as alternative to the traditional SW frame- work by replacing one-dimensional projection lines with more structured domains, known as tree systems. These sys- tems function similarly to lines but incorporate a more so- phisticated and interconnected structure. Instead of project- ing measures onto individual lines, this method distributes them across multiple linked lines, forming hierarchical structure. This approach enhances the representation of topological information while preserving the computational efficiency by leveraging the closed-form solution of OT problems in tree-metric spaces (Indyk & Thaper, 2003; Le et al., 2019). However, Tree-Sliced frameworks are still re- stricted to linear projections, the integration domains used in (Tran et al., 2025c:b) remain confined hyperplanes. In contrast, several advanced projection techniques have been investigated for SW, enhancing its flexibility and effective- ness (Kuchment, 2006; Kolouri et al., 2019; Chen et al., 2022; Bonet et al., 2023). Building on these advancements, this paper introduces a novel framework for the tree-sliced method that integrates nonlinear projections, further broad-
1. Introduction
Optimal Transport (OT) (Villani, 2008; Peyré et al., 2019) and Sliced-Wasserstenn (SW) (Rabin et al., 2011; Bonneel et al., 2015) provide geometrically meaningful metrics in the space of probability measures, making them widely applicable across various fields. These include machine
'Equal contribution 'Co-last authors 'innnniversity NNational University of Singapore Movian AI "The Institute of Sta- tistical Mathematics. Correspondence to: Viet-Hoang Tran choang.tranviett www.wiiittu.uuu.uu......
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Tree-Sliced Wasserstein Distance with Nonlinear Projection
ening its scope and improving its performance.
Wasserstein distance (W) (Villani, 2008) between 14, v vis:
Contribution. Our contributions are three-fold:
Joza d(e, v) di dr(x, y). =EP(U.)
w(
inf
(I)
We introduce the Generalized Radon Transform and Spatial Radon Transform on Systems of Lines, ex- tending previous variants that were limited to linear provide
Sliced Wasserstein dístance. The Radon Transform (Hel- gason, 2011) R (R)) - L x 1d is:
projections. Along with these extensions, theoretical results and proofs demonstrating their in- jectivity. Furthermore, we generalize the conventional Euclidean framework by proposing a spherical tree- sliced version tailored for nonlinear Radon transforms.
(1, / f(v) o(t (v. f3) dy,
(2)
where is the Dirac delta function.
The Sliced Wasserstein distance (SW) (Rabin al., 2011;
We introduce Tree-Sliced distances based on the pro- posed Radon Transform, namely the Circular Tree- Sliced Wasserstein distance and the Spatial Tree-Sliced Wasserstein distance, along with an extended version for the spherical setting, called the Spatial Spherical Tree-Sliced Wasserstein distance. Furthermore, we ex- amine different choices of functions that define the non- linear projections, analyze their computational com- plexity and explain why certain choices lead to more efficient metrics.
Bonneel al., 2015) between pr, E P(RR)) is:
{ W(Rfu(.0),
SW(H, v) = Jga-1
do),), (3)
where a u($$-1) is the uniform distribution on the sphere, and Su are the probability density functions of M,v, respectively. The one-dimensional Wasserstein distance in Eq. (3) has the closed-form W(Opp, 0tv) So (' (2) F (..0 (z)ldz, where (=))(,)))))) FRI FRIA(-.0) are the cumulative distribution func- tions of ()) ()) respectively. The Monte Carlo method is employed to approximate the intractable integral in Eq. (3):
We assess the proposed metrics across various tasks, including gradient flows and generative models, on both Euclidean and spherical data, highlighting their practical effectiveness and computational efficiency.
SW(H,)) v) = 1 Σ W(RJ((10) (4)
Organization. The structure of the paper as follows: Section 2 provides an overview of different variants of the Wasserstein distance, while Section 3 explores various forms of the Radon Transform with Nonlinear Projection used to derive corresponding Wasserstein distances. Section 4 in- troduces novel Radon Transforms on Systems of Lines with Nonlinear Projection and examines their injectivity. In Sec- tion 5, two new Tree-Sliced Wasserstein distances associated with the proposed transform are introduced, along with an analysis of their fundamental components. Finally, Section 6 assesses the performance of the proposed methods on both Euclidean and Spherical data. Additional materials related to the spherical settings of the proposed method, as well as background, theoretical foundations, and supplementary content, are provided in the Appendix.
L 1=1
where .I....,0. O2 are drawn independently from
Tree-Sliced Wasserstein distance on Systems of Lines. Rather than projecting functions onto lines, i.e., directions in gd-1, as in the original Radon Transform, (Tran et al., 2025c;b) introduces an alternative approach that replaces lines with different metric measure spaces, known as tree systems. These tree systems are formed by connecting multiple copies of real lines, creating a structured space. Functions are then partitioned using a splitting mechanism and projected onto these spaces, effectively capturing the positional information of both measure supports and slices. Further details on this Tree-Sliced framework can be found in Appendix A and Appendix C.I.
2. Preliminaries
3. Radon Transform with Nonlinear Projection
This section reviews the Wasserstein distance between mea- sures and its various sliced variants. For simplicity, the focus is on measures with a finite first moment, while measures with a finite ph moment are treated analogously.
In this section, we explore nonlinear projections utilized in various existing Sliced Wasserstein variants, and compare them with the traditional linear projection in the original Radon Transform that leads to the Sliced Wasserstein dis- tance. Based on these observations, we provide an overview of how these nonlinear projections can be incorporated into the framework of the Radon Transform on Systems of Lines.
Wasserstein distance. Given a measurable space Ω en- dowed with a metric d. Let μ, v E P(2). and P(u, 2/) be the set of distributions æ coupling between 4 and v. The
2

---

Tree-Sliced Wasserstein Distance with Nonlinear Projection
3.1. Generalized and Spatial Radon Transforms
3.2. Incorporating Nonlinear Projectional Framework into Systems of Lines Setting
Let v/ be a set of feasible parameter, the Generalized Radon Transform G (GRT) (Kolouri et al., 2019) is defined by:
Here, we provide brief overview of the Radon Transform Systems of Lines (RTSL) to highlight the potential of integrating the nonlinear approach discussed in Section 3.1. The formal construction of RTSL with notations is detailed in Appendix A and in (Tran et al., 2025c:b). Roughly speak- ing, unlike the traditional Radon Transform in Eq. (2), which projects onto one line at a time, RTSL extends this concept by simultaneously projecting onto set of interconnected multiple lines, denoted L, using a splitting mechanism:
g L' L'(R'' (R*) - L' (R x V)
s.t. (f(t, = J f(v) 6(t - -ly, 1² dy. (5)
Here, 9: Rdd x Ω - IR called the defining function of G. The GRT of E L (IB=) is the integration of f over hypersurfaces fy E R t g(y, })} for t E RR vh E V. One common choice for the defining function occurs when T = R30 x gú-1, and is defined as the circular function:
Rº: (R II (C(C) where f (R2 (ELL)
gly. r, 0) = | - -r0ll2 Vy E Rdd (r,
(r, E V.
(6)
CEL
s.t. R2f(21 +t 02)
This choice results in the Circular Radon Transform (CRT) (Kuchment, 2006). Another possible defining function, based on homogeneous polynomials of odd degree, is pre- sented in (Rouvière, 2015). However, this represents a special case of the next Radon Transform we discuss. Given a positive integer dg, the Spatial Radon Transform H (SRT) (Chen et al., 2022) is defined by:
= -/ee
((w) (()) . aly, L).. 8 (y - Ii, xi, 0.)) dy. (8)
Here, a is a continuous map from IRd x L1 to A pre- senting the splitting mechanism, and (xx, 0.) indicates 8th line in C. Eq. (8) can be interpreted as integrating the 7th portion of f. given by f a; over the hyperplanes {y E Bd (v, - - (I 0,))-
H: L' L'RR))
L' L' (Rxx80=1) 1.
s.t. H)((t,)= = [ f(v) ( o(t-00ttttt (h(v), 0)) dy, (7) zd
Remark 3.2. It is important to note that the partitioning process depends on both y E Id and the set of lines L. This splitting mechanism is absent in the traditional Radon Transform and its variants, as seen in Eqs. (2), (5), (7). This naturally leads to the question:
where h: Rdd -s Redo is an injective continuous map. The SRT of / E (RR)) is defined as the integration of ʃ over the hypersurfaces given by {y E Rd t (h(v). 8)} for t I. 0 E 5:
Can a nonlinear projectional framework, similar to GRT and SRT, be developed for RTSL?
Remark 3.1. When Q4 = g0-1, gly. 0) = (y, 0) in Eq. (5), or d = dg, hyy) = y in Eq. (7), it is clear that the GRT and SRT recover the original Radon Transform in Eq. (2).
Notably, the presence of the map a introduces trade-off between the effectiveness of the induced Wasserstein metric and the associated theoretical guarantees. Empirical find- ings indicate that, given the same number of projections (and consequently the same computational cost), the dis- tances derived from RTSL surpass those obtained from the original Radon Transform. However, incorporating shifts the transformation from operating on straight lines to more intricate space. This transition may compromise fundamen- tal properties of the transform, such as injectivity, which might no longer be assured in this new framework.
By definition, the SRT is a special case of the GRT, but these two transforms can be interpreted from two distinct perspectives. For general function g. the GRT represents a projection along hypersurfaces in 12d, detined by the level sets of g. In contrast, for general function h, the SRT involves mapping functions from IR to new space dR where the Radon Transform is then applied.
Well-definedness and injectivity. Certain conditions on g in GRT and h in SRT are necessary to ensure the well- definedness of GRT and SRT. Additionally, since injectivity typically required for Radon Transform variants. spe- cific assumptions on g and h are made to achieve this prop- erty. However, since these properties, along with the in- verse problem related to Radon Transform variants, remain long-standing research questions (Beylkin, 1984; Ehren- preis, 2003; Uhlmann, 2003; Homan & Zhou, 2017) and fall beyond the scope of this paper, we restrict our discussion to mentioned examples of g and h found in the literature.
In the next sections, we propose an approach for incorporat- ing the nonlinear projectional framework into the systems of lines setting, thereby generalizing the current RTSL and inducing new variants of the tree-sliced distance.
4. Radon Transform on Systems of Lines with Nonlinear Projection
In this section, we extend the current Radon Transform on Systems of Lines (Tran al., 2025c;b) by incorporating
3

---

Tree-Sliced Wasserstein Distance with Nonlinear Projection
nonlinear projections. We then analyze key properties of the resulting transforms, including injectivity.
Injectivity. For injectivity of CR and H we refer to the concept of F(d)-invariance in splitting maps intro- duced in (Tran et al., 2025b). The group E(d) represents the Euclidean group, which consists of all transformations of Fuclidean space Rd that preserve the Fuclidean distance between any two points. Through the canonical action of B(d) on R4, an induced action of E(d) on Lg follows. Ap- pendix A provides a formal description of the underlying group actions associated with these equivariant construc- tions. A splitting map E C(IR x LA AR-1) is E(d)- invariant if:
4.1. Nonlinear Radon Transform on Systems of Lines
Given a positive integer k representing the number of lines in tree system, and continuous splitting map function E C(RR x Lig 1k-1) defining the splitting mechanism. Let C be a system of k lines in L and scalar r' > 0. For a function / E L.a (IR.)) define the function CR2, E (C)
as follows:
CR2 + 0.))
algy.gL) = e(y,.))
(14)
= =/6- f(y) aly, L)2 6 (t- 11y 1ly I;- r0:112) dy. (9)
for all (v. C) E R x Lg and E E(d). We have two results about injectivity of operator CR and H
The Circular Radon Transform on Systems of Lines (CRTSL) is defined as the operator:
Theorem 4.2. For an Bnda-invariant splitting map a C(R x Lg. Ak-1), CR is injective.
CR (RR)
II (2)
Theorem 4.3. For an E(da) (dd)--invvrian splitting map e C(R x L Ax-1). H. is injective.
f
(CR2,rf) 2er (10)
This is analogous the CRT described in Section 3.1. In the case of SRT, consider positive integer dg, an injective continuous map h: Rd - Rdo, and splitting map E C(R x Lg Ax-1). Let L be a system of lines in L For a function f E I. (R.)), define the function H2 as:
The proof' of Theorem 4.2 and Theorem 4.3 are provided in Appendices B.3 and B.4. Intuitively, the reason why i(d)-invariance in splitting maps ensures the injectivity of CR stems from the fact that the circular defining function in Eq. (6) is primarily based on the Euclidean norm || ll2, which itself is an E(d) function. Similarly, the reason why E(do) innvariance in splitting maps guarantees the injectivity of H is that this property is essential for achieving injec- tivity in the standard Radon Transform Systems of Lines in Rd as discussed in (Tran et al., 2025b).
Hz + t- t.))) = -
(11)
Rd ((v) ((h()())))))))))))) (1), L): 8 (h(y) - ns, 0.)) dy.
The Spatial Radon Transform on Systems of Lines (SRTSL) is defined as the operator:
Ha : (R()))
II (2)
4.3. Spatial Spherical Radon Transform on Spherical Trees
DEL
(H2f)
(12)
In (Tran et al., 2025a), the tree-sliced framework is extended to functions defined on hyperspheres. The techniques pre- sented in Sections 4.1 and 4.2 can be adapted to the spher- ical setting. We provide a brief derivation here, while a more detailed background and notation are given Ap- pendix C.1. For f E L (84), we recall the Spherical Radon Transform on Spherical Trees (SRTST), which transforms f to R7f E L LT),,, where:
Remark 4.1. It is important to note that when the system of lines consists of a single line, iee k = 1, CR" and H recover the GRT and the SRT, respectively.
Properties of CR and 4 are discussed in the next part.
4.2. Well-definedness and Injectivity
Well-definedness. Given the setting of CR and Ha as defined in Eqs. (9), (11), we have CRE, E (C) and H2f E (C) for f E L (IRd). Furthermore, we have the bounds:
(t,((trg))
f()) aly. T). arccos (a. dy, (15) f gg
Since the literature on defining functions for GRT on the sphere is limited, we focus only on the spatial version of SRTST. Given a positive integer dg and an injective con- tinuous map h: gd + gdo, the Spatial Spherical Radon Transform on Spherical Trees transforms f E T T,*(S)) to
|CRR,,flle < lifll and ||Hfflcc < Iffll. (13) The proofs for these properties are provided in Appen- dices B.1 and B.2. Additionally, these proofs imply that if P(RR),, then ,R H& E P(C).
4

---

Tree-Sliced Wasserstein Distance with Nonlinear Projection
H4.5 (T), where: H95(t, (()())) T)- 8(t arccos (z, h(y))) dy- g4
Theorem 5.3. CircularTSW and SpatialTSW are metrics the space P(
(16)
The proof of Theorem 5.3 is presented in Appendix B.5. The algorithms for CircularTSW and SpatialTSW presented in Appendix D.1 (Alg. and respectively).
As stated in (Tran et al., 2025a), Odd + i)-ivaariance is necessary property of the splitting map to ensure the injectivity of R A similar result holds in our setting, as described below.
5.2. Components in CireularTSW and SpatialTSW
Both CircularTSW and SpatialTSW distances depend on the choice of splitting maps a. Additionally, CircularTSW is influenced by the parameter r > 0, whereas SpatialTSW is determined by the selection of the injective map h.
Theorem 4.4. For an (dg + -iiinaarinnt splitting map C( x T 11-1), Hº injective.
The detailed derivation of the Spatial Spherical Radon Trans- form on Spherical Trees and the proof of Theorem 4.4 are provided in Appendices C.2 and C.3.
Splitting maps a. For both CircularTSW and SpatialTSW, we follow the construction of the splitting map a as pro- posed (Tran al., 2025b), defined as:
5. Tree-Sliced Wasserstein Distance with Nonlinear Projection
a(x, C)2 = softmax ())=))
(19)
In this section, we propose new distance between measures derived from the variants of the Radon Transform introduced Section 4. We also examine different choices of functions that define the nonlinear projections and explain why certain choices lead more efficient metrics.
where d(x, C): represents the distance between and 7t line of L. This choice of a ensures the E(d))invariance while also incorporating positional information between point and the tree system. Consequently, it results in meaningful and varied mass distributions that adapt to each specific system. Moreover, the use of the softmax function guarantees that a produces a valid probability vector in the standard simplex A The splitting map for CircularTSW is further discussed in Appendix D.4.
5.1. Definition of Tree-Sliced Distances
For two probability measures j and with density function Jn and fu. and a fixed > O, the Circular Tree-Sliced Wasserstein Distance (CireularTSW) between μ and v is defined as the average Wasserstein distance on the tree- metric space L between the CRTSL of fu and fu. Following (Tran et al., 2025c:b), this averaging is taken over the space of trees Tí C L.g, according to a distribution a on Ty which arises from the tree sampling process.
Radius r in Circular In the formulation of the orig- inal Circular Radon Transform, the radius plays a crucial role. Together with E ga-1, the term r spans the entire space Rd. This ensures that the level sets defined by the defining circular function in Eq. (6) can represent arbitrary (d diimessionl spheres in Rd:
Definition 5.1. The Circular Tree-Sliced Wasserstein Dis- tance between j and v in P(Rd)) is defined by:
{y E Rd 1= 11y r0112] for L 1>
(20)
CircularTSW(H.t)
However, in the framework of Tree-Sliced methods, since the sources within tree systems already encompass the en- tire space Fd, considering all values of r O becomes redundant. In other words, for fixed r the level sets arising in the Circular Radon Transform on Systems of Lines from Eq. (9) can still effectively represent arbitrary (d - 1) dimmensiona spheres in Rdd
= FICR (L). (17)
Similarly, given choice of the continuous injective map h: d Rd in SRTSL, we have the definition of the Spatial Tree-Sliced Wasserstein Distance (SpatialTSW) be- tween μ and v.
Definition 5.2. The Spatial Tree-Sliced Wasserstein Dis- tance between l and v in P(R.d) defined by:
{v E Rddd t Bly I, r0,112} for to 0.
(21)
This is why, in the definition of CircularTSW, specific r > 0. Furthermore, we want to examine CircularTSW
fix
SpatialTSW v),
Jupp Heff (C). (18)
special case of CircularTSW when r = O. In this scenario, by selecting the concurrent tree space as proposed in Tran et al. (2025b), the practical im- plementation of CircularTSW becomes significantly more
Both CireularTSW and SpW distances are, indeed, metrics on the space P(RR)) of measures on Rd.
5

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
5.3. Computational Complexity
Consider two discrete measures μ, P(RR)) with m, support points, respectively. The computational complexity of the original Sliced Wasserstein distance is OLLn log n + Ldn), where L represents the number of samples used in the Monte Carlo approximation (Peyré et al., 2019; Nguyen et al., 2024a). In comparison, the computational complex- ity of the Tree-Sliced Wasserstein (TSW) distances, such as TSW-SL in (Tran et al., 2025c) or Db-TSW in (Tran et al., 2025b), is O (LLn log n + Lkdn), where I. denotes the number of tree samples used in the Monte Carlo ap- proximation, and k is the number of lines per tree. This computational difference highlights why fair compari- son between the sliced method and the tree-sliced method requires ensuring that the total number of directions re- mains the same. For SpatialTSW, its computational com- plexity is O log n + Lkdgn), with additional ini- tial cost for computing the function h. This complex- ity matches TSW-SL and Db-TSW but operates in the transformed space Rde. For CircularTSW with a general e O, the complexity remains ( log n + Lkdgn), yet it achieves faster empirical runtime than DD-TSW by computing vector norms instead of vector products. No- tably, for r to OLn log + Lkdgn). This reduction arises because the O(n log n) sorting step per line is required for only single line, rather than all k lines in a tree, as illustrated in Figure 1. The empirical efficiency of CircularTSW and CireularTSW,= is demonstrated in Figure 2, where we use I. 10000 for SW and I. 2500,
6
An illustration depicting CircularTSW,-d Given a support, the projection coordinates are identical when projected onto k lines.
Figure I:
efficient, reducing computational complexity. In summary, transforming a measure μ E P(R") involves projecting its support onto k lines in the tree system. When r = O, all sup- port points of share the same coordinates when projected onto these 4 lines. As result, projecting and sorting cost is reduced. Figure I illustrates this phenomenon. Furthermore, emphasize that CircularTSW. r is specifically designed for tree settings, where splitting maps play crucial role. Since the coordinates are identical across these lines, the tree structure and the distance-based splitting map are nec- essary to distinguish between the lines. Empirical results in Appendix D.55 show that CireularTSW r tree setting but poortly in original sliced setting.
CireularTSW improves complexity
r= performs well
for Tree-Sliced methods, following the practical setting used in the Diffu- sion Model experiment. CireularTSW w- scales efficiently with the number of supports and is the only Tree-Sliced method that closely matches the speed of vanilla SW.
The choice of the map h in SpatialTSW. As discussed in Section 4, the map h in SpW must be both in- jective and continuous. One approach to selecting this map based on odd degree homogeneous polynomials. How- ever, following the constructions in (Kolouri et al., 2019; Rouvière, 2015) can lead excessively large new di- mension dg (d)). To address this, propose alternative approach: Consider h Rd + Rd defined by:
5.4. Spatial Spherical Tree-Sliced Wasserstein Distance
From the Spatial Spherical Radon Transform on Spherical Trees, we introduce a spherical variant of the Spatial Tree- Sliced Wasserstein distance, referred to as SpatialSTSW distance. A detailed derivation, along with the selection of the corresponding injective map and theoretical proofs for the SpatialSTSW distance, is provided Appendix C.4.
I(I1,..., (Z1()............ .... Ja((a))),,,,((
(22)
where (fs: IR - R are injective and continuous func- tions. A simple choice for fi is an odd-degree polynomial that remains injective, such as f((x) = zi + x.. Another ap- proach, inspired by (Chen et al., 2022), involves concatenat- ing the input with the output of a neural network by concate- nating input with an arbitrary neural network o(-). Specif- ically, we define h: Rdd Rd+d' as h(x)) = (a,, ((x)), where Φ: Rd + gdd is neural network. This choice in- troduces learnable parameters for h, offering a trade-off between potentially improving performance and increasing computational cost.
6. Experimental Results
In this section, we thoroughly assess the validity of our proposed metric through extensive numerical experiments on both Euclidean and spherical datasets.
6.1. Euclidean Datasets
Denoising Diffusion Generative Adversarial Network. This experiment investigates training denoising diffusion
6

---

Tree-Sliced Wasserstein Distance with Nonlinear Projection
Table I: Fréchet Inception Distance (FID) scores and per-
epoch training times of different DDGAN variants for un- conditional generation on CIFAR-1O.
SN
arTS
Te TWW
Model
FID J Time/Epnch(s) 1.
Db-TOW TFW-G G.
DDGAN (Kiao et al., 2021)
3.64
188
Puntme (113)
30-
SW-DD (Nguyen al., 2024h)
2.90
192
DSW-DD (Nguyen et al., 2024b)
2.88
1268
EBSW-DD (Nguyen al., 2024b)
2.87
188
RPSW-DD (Nguyen al., 2024b)
2.82
194
IWRPSW-DD (Nguyen et al., 2024b)
2.70
194
20
TSW-SL-DD (Tran et al., 2025c)
2.83
249
Db-TSW-DD (Tran et al., 2025b)
2.60
256
190
403
Db-TSW-DD (Tran et al., 2025b)
2.53
262
Number of Supports In]
SpatialT (ours)
2.52
262
CircularTSW-DD (ours)
2.33
234
Figure 2: Runtime comparison of our proposed distances and baseline methods. We randomize the measures and projection directions and benchmark runtime over 10 runs. To is comparable to SW and significantly out- performs existing Tree-Sliced distances in terms of speed.
CircularTSW -oDD (nurs)
2.48
211
the gradient of the distance function 'D with respect to par We evaluate SpatialTSW, CircularTSW, CircularTSW, and several established Sliced- Wasserstein (SW) variants, including vanilla SW (Bonneel et al., 2015), MaxSW (Desh- pande al., 2019), LCVSW (Nguyen & Ho, 2023), SWGG (Mahey al., 2023), alongside the recently introduced Tree- Sliced distances, such as TSW-SL (Tran et al., 2025c), Db- TSW-, and Db-TSW (Tran et al., 2025b). We conduct experiments on the 25 Gaussians dataset and use the Wasser- stein distance to evaluate the average distance between the source and target distributions. We report results over runs at iterations 500, 1000, 1500, 2000, and 2500.
models for unconditional image synthesis. Following the approach of Nguyen et al. (2024b), we incorporate a Wasser- stein distance into the Augmented Generalized Mini-batch Energy (AGME) loss function of the Denoising Diffusion Generative Adversarial Network (DDGAN) (Xiao et al., 2021). We benchmark our proposed methods SpatialTSW- DD, Cir,ularTSW, and CireularTSW - against Sliced and Tree-Sliced Wasserstein-basee DDGAN variants, detailed in Table 1. All models are trained for 1800 epochs on the CIFAR10 dataset (Krizhevsky et al., 2009). For vanilla SW and its variants, we follow the parameter set- tings from Nguyen et al. (2024b), using L = 10000. For Tree-Sliced methods, including our own, we adopt the con- figuration from Tran al. (2025b), setting 2500 and k == 4. Further details on this experiment can be found in Appendix D.6
The results presented in Table 2 demonstrate that Spa- tilTTSW achieves the best performance across all itera- tions, reaching a final W2 distance of 1.17e-7 at the last step. This represents a significant improvement over vanilla SW (3.59c-2) 2) and LCVSW (9.28c-33 3). Furthermore, com- pared to the best existing Tree-Sliced distances, SpatialTSW achieves better results than Db-TSW (1.3e-7), 7), exhibiting faster convergence, and maintaining similar computational efficiency. In this experiment, SpatialTSW outperforms CircularTSW, aligning with the findings from (Kolouri et al., 2019), where polynomial-base based defining functions yield superior results over circular defining functions in gradient flow tasks. Nevertheless, both CircularTSWW and CireularTSW, still achieve better results than vanilla Sw while offering significant computational speedups. Specifi- cally, CireularTSW and CircularTSW. == are approximately 5% and 16% faster than vanilla SW, respectively.
The results in Table I show that SpatialTSW-DD, CircularTSW-DD, and CircularTSW,... DDD achieve no- table improvements in FID compared to all baselines. They surpass the current state-of-the-art OT-based DDGAN, Db- TSW-DD- (Tran et al., 2025b), by margins of 0.01, 0.2, and 0.05, respectively. Additionally, our methods offer faster training times compared to existing SriceS approaches. CircularTSW-DD and CircularTSW reduce training time relative to Db-TS -D-DSW-DD- by 10% and 19%, respectively. These enhancements in both training efficiency and model performance underscore the practical advantages of our pro- posed methods.
6.2. Spherical Datasets
Gradient Flow on The Sphere. We now evaluate the abil- ity to learn distributions by iteratively minimizing d(v, p),, where d is a distance metric such as SSW (Bonet et al., 2022), S3W variants (Tran et al., 2024a) and STSW (Tran et al., 2025a). In line with previous works (Tran et al., 2024a; 2025a), we consider a mixture of 12 von Mises- Fisher distri-
Gradient Flow. The goal of gradient flow is to minimize the distance between a source distribution μ and target dis- tribution v through gradient- based optimization. The update rule follows 0,44 = - D(4,, v), Mo = _N 1), where 14 represents the distribution at time f, and V'((u, v) is
7

---

Tree-Sliced Wasserstein Distance with Nonlinear Projection
Table 2: Average Wasserstein distance between source and target distributions of 5 runs on 25 Gaussians datasets. All methods use 100 projecting directions.
Table 4: Accuracy of the linear classifier on encoded (E) features and projected (P) features on S ARI-S3W and RI-S3W use rotations.
Methads
licration
Time/len(a)
Method
Acc. R(%) 1
Acc. P(%)
Time (s/ep.)
500
1000
1500
2000
2500
Hypersphere
79.78
74.60
13.10
SW
4.21e-1
1.54e-1
7.72e-2
4.97e-2
3.59e-2
0.0018
MasSW
5.23e-1
2.36e-1
1.236-1
8.04-22
6.76e-2
SimCLR
79.86
72.79
12.71
SWOG
6.59c-1
3,62c-
199c-1
9.07e-2
4.420-2
0.0019
LCVSW
3.46e-1
2
2.26e-2
1.3le22
9,28c:3
0.0019
SSW
70.37
64.76
13.31
TSW-LL
3,49--1
8.100-2
1.06-22
2.68e-3
3.16c-6
0.0019
S3W
78.53
73.73
12.90
D6-TSW
3,50c-1
8.12c-2
1009-22
1.77e33
1,30=7
0.0020
RI-S3W (5)
79.96
74.02
13.08
DS-
3.52e-1
7.69e-2
2.73e-2
2.56e-3
2.03e/6
0.0021
ARI-S3W (5)
80.06
75.10
13.01
SpanialTS
3.20e-1
3.44e-2
2.95e-3
3,97e-4
1.17-77
0.0021
Cir
4.28c-1
1.200-1
3.48-22
1.41-22
7.86e-3
0.0017
STSW
80.51
76.79
12.81
CwealarTSW a
4.32e-1
1,23o-1
3.41o-2
1/45-22
8.94p-3
0.0015
SpatialSTSW
80.68
77.31
12.87
Table 3: Average Log of the Wasserstein distance between
source and target distributions over 5 runs on a mixture of 12 vMFs.
ing a similar approach to Bonet et al. (2022); Tran et al. (2024a; 2025a), we use the above objective function to pre- train ResNet18 (He et al., 2016) encoder on CIFAR-10 (Krizhevsky et al., 2009) for 200 epochs and then train a linear classifier to evaluate learned features. The results in
Methods
Epoch
Time/Epoch(a)
100
150
SSW
-2.1274
-2.7893
-2.9226
-2.9882
-3.0313
0.4323
S3W
-2,0204
-2.1920
-2.2615
-2.2699
-2,2734
0.0151
-2.1107
-2.5163
-2.7295
-2.8568
-2.9447
0,0182
R1-S3W (5)
-2,4399
-2,8273
-3,0093
-3.1234
-3.2145
0,0503
Table 4 indicate that SpatialSTSW achieves the best per-
ARLS3W (30)
-2,6508
.3.0279
.3.2405
-3,3855
.3.661
0,1884
formance compared to various baseline methods, including Hypersphere (Wang & Isola, 2020), SimCLR (Chen et al., 2020), SSW, S3W variants, and STSW.
STSW
-2.9545
-3.3322
.3.992
..3623
4.6486
SpatialSSS
.2.8224
-3,4226
-4.0903
-4.5368
-4,6599
0.0145
butions (vMFs) from which we have access to a sample set M { with M/ = 2400. The optimization procedure uses projected gradient descent (Bonet et al., 2022), applied on the sphere with full-batch size. Table 3 illustrates the evo- lution of the log 2- Wasserstein distance at epochs 50, 100, 150, 200, and 250, averaged over 5 runs. From these results, SpatialSTSW demonstrates better performance compared to baselines while maintaining computational efficiency close to STSW.
Sliced- Wasserstein Auto-Encoder. In this study, we uti- lize the work introduced by Kolouri et al. (2018) to evaluate the per- formance of various distances, including SW, SSW (Bonet et al., 2022), S3W variants (Tran et al., 2024a), and STSW (Tran et al., 2025a). The target of SWAE is to ensure that the encoded embeddings follow a predefined prior distri- bution q in the latent space. Let the encoder be denoted as op : X' gd and the decoder as gd X. The optimizing objective is defined as:
Auto-Encoder (SWAE) frame-
Self-Supervised Learning (SSL.. In earlier work, Wang & Isola (2020) have shown that the contrastive objective can be broken down into alignment loss, which ensures that representations of similar inputs are close together, and uniformity loss, which prevents representations from col- lapsing by encouraging them to spread out evenly. Inspired by Bonet et al. (2022), we replace the Gaussian kernel in uniformity loss with our SpatialSTSW:
min Eawp Je(æ, (+((=))) + X SpatialSTSW (P:P. q) P2.0
where p represents the data distribution, À acts as a regular- ization weight, and c(:.-) measures the reconstruction error. For reconstruction loss, we use Binary Cross Entropy (BCE) and adopt mixture of 10 von Mises-Fisher (vMF) distri- butions the prior. We report results in Table 5, where we evaluate performance using the same metrics as in Tran et al. (2024a; 2025a). We observe that SpatialSTSW achieves better results in terms of log W2 and NLL while maintain- ing a competitive reconstruction loss (BCE) and efficient computation times.
n L
Σ n
Aligament loss
X (SpatialSTSW<>> ,D) SpatialSTSW (n),
2
Uniformity loss
7. Conclusion
where V = U($d) is the uniform distribution on Sd, A B E (d)(d+1) are feature embeddings of two aug- mented views of the same image and > 0 is regular- ization factor that balances the loss components. Follow-
This paper introduces the Circular Tree-Sliced Wasserstein Distance (CircularTSW) and the Spatial Tree-Sliced Wasser- stein Distance (SpatialTSW) as novel approaches for com- paring probability measures in Euclidean spaces. These
8

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
Table 5: CIFAR-I0 results for SWAE evaluated on latent
bolic Wasserstenn via geodesic and horospherical projections. In Topological, Algebraic and Geometric Learning Workshops 2023, pp. 334-370. PMLR, 2023.
regularization.
Method
log W2
NLL
BCB
Time (s/ep.)
SW
-3.3181
-0.0010
0,6330
6,8939
Bonneel, N., Rabin, J., Peyré, G., and Pfister, H. Sliced and
SSW
-2.3425
0.0037
0.6316
16.8639
Radon Wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision, 51:22-45, 11,22-45,2015 2015.
S3W
-3.3181
0.0018
0.6307
8.9703
RI-S3W
-3.1857
-0.0034
0.6357
10.1904
ARI-S3W
-3.3850
0.0020
0.6328
9.5882
Bunne, C., Papaxanthos, L., Krause, A., and Cuturi, M.
STSW
-3,4098
-0.0045
0,6347
7.1623
Proximal optimal transport modeling of population dy-
SpatialSTSW
-3,4254
-0.0049
0.6368
7.2811
namics. In International Conference Artificial Intelli- gence and Statistics, pp. 651 6511-6528. PMLR, 2022.
approaches integrate nonlinear projection techniques from the Sliced Wasserstein distance into the recent Tree-Sliced Wasserstein framework, resulting in enhanced performance and more efficient metric computations. The paper presents formal derivation of these metrics and provides compre- hensive theoretical guarantees to ensure their practical ap- plicability. Furthermore, the proposed techniques are ex- tended to the tree-sliced framework with a spherical set- ting. Experimental evaluations show that CircularTSWW Spa- tialTSW, and their spherical variant consistently outperform state-of-the-ar Sliced Wasserstein and Tree-Sliced Wasser- stein methods across various tasks, including gradient flows and diffusion models, while maintaining comparable improved runtime efficiency. These results highlight Tree- Sliced Wasserstein distance as a promising and impactful research direction, complementing the Sliced Wasserstein distance in practical applications.
Chapel, L. and Tavenard, R. One for all and all for one:
Efficient computation of partial wasserstein distances on the line. In International Conference on Learning Repre- sentations, 2025.
Chapel, L., Tavenard, R., and Vaiter, S. Differentiable generalized sliced wasserstein plans. arXiv preprint 2505::2500.22455 2025.
Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of visual rep- resentations. In International conference on machine learning. pp. 1597-1607. PMLR, 2020.
Chen, X., Yang. Y., and Li, Y. Augmented sliced Wasser- stein distances. In International Conference on Learning Representations, 2022. URL https: httpp:://pppreeeeeeeeeee /opnrevien.....ee id- htt.:forrmmiid-iiggggggnan iMgTLyEWNO.
Acknowledgements
Deshpande, I., Hu, Y.-T., Sun, R., Pyrros, A., Siddiqui, N., Koyejo, S., Zhao, Z., Forsyth, D., and Schwing, A. G. Max-sliced Wasserstein distance and its use for GANs. Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10648-10656, 2019.
We thank the area chairs and anonymous reviewers for their comments. TL gratefully acknowledges the support of JSPS KAKENHI Grant number 23K 1243, and Mitsui Knowl- edge Industry Co., Ltd. grant.
Ehrenpreis, L.
The universality of the Radon transform. OUP Oxford, 2003.
Impact Statement
This paper presents work whose goal is advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Fan, J., Haasler, I., Karlsson, J., and Chen, Y. On the complexity of the optimal transport problem with graph- structured cost. In International Conference on Artificial Intelligence and Statistics, pp. 9147-9165. PMLR, 2022.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
References
Beylkin, G. The inversion problem and applications of the generalized Radon transform. Communications on pure and applied mathematics, 37(5):579-59,, 1984.
Helgason, S. The Radon transform on R Integral Geome- try and Radon Transforms, 1-62, 2011.
Bonet, C., Berg, P., Courty, N., Septier, F., Drumetz, L., and Pham, M.-T. Spherical Wasserstiss...iii.. arXiv preprint arXiv:2206.08780 2022.
Ho, J., Jain, A., and Abbeel, Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:68400668 6851, 2020.
Bonet, C., Chapel, L.,, Drumetz, L.,, and Courty, N. Hyper-

---

Tree-Sliced Wasserstein Distance with Nonlinear Projection
Homan, A. and Zhou, H. Injectivity and stability for generic class of generalized Radon transforms. The Jour- nal of Geometric Analysis, 27:1515-1529, 2017.
Le, T., Nguyen, T., and Fukumizu, K. Optimal transport for measures with noisy metrie. In International Con- ference on Artificial Intelligence and Statistics. PMLR, 2024b.
Hua, X., Nguyen, T., Le, T., Blanchet, J., and Nguyen, V. A. Dynamic flows on curved space generated by labeled data. In Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence, LJCAI-23, pp. 3803-3811 11, 2023.
Liu, L., Pal, S., and Harchaoui, Z. Entropy regularized optimal transport independence criterion. In International Conference on Artificial Intelligence and Statistics, pp. 11247-11279. PMLR, 2022.
Indyk, P. and Thaper, N. Fast image retrieval via embed- dings. In International workshop on statistical and com- putational theories of vision, volume 2. pp. 5, 2003.
Luong, M., Nguyen, K., Ho, N., Haf, R., Phung, D., and Qu, L. Revisiting deep audio-text retrieval through the lens of transportation. XrXiv preprint arXiv:2405. 10084, 2024.
Mahey, G., Chapel, L., Gasso, G., Bonet, C., and Courty, N. Fast optimal transport through sliced generalized Wasser- stein geodesics. In Thirly-seventh Conference on Neural Information Processing Systems, 2023. URL https: / openreview net/ forum? fetimmid==n333 3d=n3xuddvhN.
Kessler, S., Le, T., and Nguyen, V. SAVA: Scalable learning- agnostic data valuation. In The Thirteenth International Conference on Learning Representations, 2025.
Kinga, D., Adam, J. B., et al. A method for stochastic optimization. In International conference on learning representations (ICLR), volume 5, pp. 6. San Diego, California;, 2015.
Mena, G. and Niles-Weed, J. Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem. In Advances in Neural Information Processing Systems, pp. 4541-4551, 2019.
Kolouri, S., Rohde, G. K., and Hoffmann, H. Sliced Wasser- stein distance for learning Gaussian mixture models. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp.
Nadjahi, K., Durmus, A., Jacob, P. E., Badeau, R., and Sim- sekli, U. Fast approximation of the sliced-Was Wasserstein distance using concentration of random projections. Ad- vances in Neural Information Processing Systems, 34: 12411-12424, 2021.
2018.
Kolouri, S., Nadjahi, K., Simsekli, U., Badeau, R., and Rohde, G. Generalized sliced Wasserstein distances. Advances in neural information processing systems, 32, 2019.
Nguyen, K. and Ho, N. Sliced Wasserstein estimation with control variates. arXiv preprint 2tt.::22230000002.,. 2023.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. 2009.
Nguyen, K., Ho, N., Pham, T., and Bui, H. Distributional sliced-Wii Wasserstein and applications to generative model- ing. arXiv preprint arXiv:2002.07367, 2020.
Kuchment, P. Generalized transforms of Radon type and their applications. Proceedings of Symposia in Applied Mathematics, 63, 01 2006. doi: /0.1099//sapm/033/ 2208237.
Nguyen, K., Bariletto, N., and Ho, N. Quasi-Monte Carlo for 3D sliced Wasserstein. In The Twelfth In- ternational Conference on Learning Representations, 2024a. URL htth:t::: http///ppppp/eecpe////p/p httpeeiwwwwwweeeeeeeeeeeeeeeeeee forum? id=Md47f7HEX
/
Lavenant, H., Claici, S., Chien, E.,, and Solomon, J. Dynam- ical optimal transport on discrete surfaces. In SIGGRAPHI Asia 2018 Technical Papers. pp. 250. ACM, 2018
Nguyen, K., Zhang, S., Le, T., and Ho, N. Sliced Wasser- stein with random-path projecting directions. arXiv preprint arXiv:2401.15888 2024b.
Le. T., Yamada, M., Fukumizu, K., and Cuturi, M. Tree- sliced variants of Wasserstein distances. Advances in neural information processing systems, 32, 2019.
Nguyen, T., Pham, Q.H.,, Le, T., Pham, T., Ho, N., and Hua, B.-S. Point-set distances for learning representations of 3D point clouds. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 10478-10487, 2021a.
Le, T., Nguyen, T., Phung, D., and Nguyen, V. A. Sobolev transport: A scalable metric for probability measures with graph metrics. In International Conference on Artificial Intelligence and Statistics, pp. 9844-9868. PMLR, 2022.
Nguyen, T. D., Trippe, B. L.,, and Broderick, T. Many pro- cessors, little time: MCMC for partitions via optimal transport couplings. In International Conference on Arti- ficial Intelligence and Statistics, pp. 3483-3514. PMLR, 2022.
Le, T., Nguyen, T., and Fukumizu, K. Generalized Sobolev transport for probability measures on a graph. In Forty- first International Conference on Machine Learning, 2024a.
10

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
Nguyen, V., Le, T., Yamada, M., and Osborne, M. A. Op- timal transport kernels for sequential and parallel neural architecture search. In International Conference on Ma- chine Learning, pp.
Tran, H., Vo, T., Huu, T., Nguyen The A., and Nguyen, T. Monomial matrix group equivariant neural functional networks. In Advances in Neural Information Process- ing Systems, volume 37, pp. 186288486608 ciates, Inc., 2024b.
PMLR, 2021b.
Curran Asso-
Nietert, S., Goldfeld, Z., and Cummings, R. Outlier-robust optimal transport: Duality, structure, and statistical analy- sis. In International Conference on Artificial Intelligence and Statistics, pp. 11691-11770 PMLR, 2022.
Tran, H. V.. Chu, T., N.gNggn--Nnn Nhat,, M.-K.,, Pham, H. T., Le, T., and Nguyen, T. M. Spherical tree-sliced Wasserstein distance. In The Thirteenth International Conference on Learning Representations, 2025a. URL https:// openreview net forum Peorummid-PPae Pid-FPOZXXE9NN.
Park, J., Lee, J., and Sohn, K. Bridging vision and lan- guage spaces with assignment prediction. arXiv preprint arXiv:2404.09632, 2024.
Tran, H. V., Nguyen- M.-K., Pham, H. T., Chu, T., Le, T., and Nguyen, T. M. Distance-based tree-sliced Wasser- stein distance. In The Thirteenth International Confer- ence Learning Representutions, 2025b. URL https: // //openreview. net/ forum? wtt.immidddiiiiiiioeeeeeeo
Peyré, G., Cuturi, M., et al. Computational optimal trans- port: With applications to data science. Foundations and Trends@ in Machine Learning, 116-6))335--607 2019.
Tran, V.-H., Vo, T. N., Huu, T. T., and Nguyen, T. M. A clifford algebraic approach to (n)-equivariant high-order graph neural networks. arXiv preprint axXiv:2410.04692, 2024c.
Pham, T., Shimizu, S., Hino, H., and Le, T. Scalable coun- terfactual distribution estimation in multivariate causal models. In Conference on Causal Learning and Reason- ing (CLeaR), 2024.
Tran, V.-H., Vo, T. N., The, A. N., Huu, T., Nguyen- Nhat, M.-K., Tran, T., Pham, D.-T., and Nguyen, T. M. Equivariant neural functional networks for transformers. arXiv preprint arXiv:24 (0..:22410002209 2024d.
Rabin, J., Peyré, G., Delon, J., and Bernot, M. Wasserstein barycenter and its application to texture mixing. In In- ternational Conference on Scale Space and Variational Methods in Computer Vision, pp. 435-446, 2011.
Tran, V.H.,, Pham, T., Tran, T., Nguyen, K., Chu, T, Le.
Rouvière, F. Nonlinear Radon and Fourier transforms. 2015. URL https: //ap:://api... http:/app..a.m.e icchooii oxc//CozpusD::2116655
T., and Nguyen, T. M. Tree-sliced Wasserstein distance: A geometrie perspective. In Forty-second International Conference on Machine Learning, 2025c. URL https: / //openreview.r net/ forum? St 9n4
Saleh, M., Wu, S.-C., Cosmo, I.., Navab, N., Busam, B., and Tombari, F. Bending graphs: Hicrarchical shape matching using gated optimal transport. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 11757-11767, 2022.
Uhlmann, G. Inside out: inverse problems and applications,
volume 47, Cambridge University Press, 2003.
Villani, C. Optimal Transport: Old and New, volume 338.
Springer Science & Business Media, 2008
Salimans, T., Zhang, H., Radford, A., and Metaxas, D. Im- proving GANs using optimal transport. arXiv preprint 1XXi::803300573. 2018
Vo, T. N., Tran, V.-H., Huu, T. T., The, A. N., Tran, T., Nguyen-Nhat, M.-K., Pham, D.-T., and Nguyen, T. M. Equivariant polynomial functional networks. arXiv preprint 140.4:2410.04213, 2024.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International conference on machine learning, pp. 2256-2265. PMLR, 2015.
Vu, A.- N., Do, T.-T., Nguyen, V.-T., Le, T., Tran, M.-T., and Nguyen, T. V. Few-shot object detection via synthetic features with optimal transport. Computer Vision and Image Understanding (CVIU), pp. 104350, 2025.
Solomon, J., De Goes, F., Peyré, G., Cuturi, M., Butscher, A., Nguyen, A., Du, T., and Guibas, L. Convolutional Wasserstein distances: Efficient optimal transportation geometric domains. ACM Transactions Graphics (TOG), 34(4):66, 2015.
Wang. J., Gao, R., and Xie, Y. Two-sample test with kernel projected Wasserstein distance. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151, pp. 8022-8055. PMLR, 2022.
Tran, H., Bai, Y., Kothapalli, A., Shahbazi, A., Liu, X., Martin, R. P. D., and Kolouri, S. Stereographic spherical sliced Wasserstein distances. In Forty-first International Conference on Machine Learning, 2024a.
Wang, T. and Isola, P. Understanding contrastive represen- tation learning through alignment and uniformity on the hypersphere. In International conference on machine learning, pp. 992999999999 PMLR, 2020.

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
Weed, J. and Berthet, Q. Estimation of smooth densities in Wasserstein distance. In Proceedings of the Thirty- Second Conference on Learning Theory, volume 99, pp. 1118-3119, 2019.
Xiao, Z., Kreis, K., and Vahdat, A. Tackling the generative learning trilemma with denoising diffusion GANs. arXiv preprint arXiv:2112.07800 2021.
Yamada, M., Takezawa, Y., Sato, R., Bao, H., Kozareva, Z.., and Ravi, S. Approximating 1-Wasserstein distance with trees. Transactions on Machine Learning Research, 2022.
ISSN 2835-8856.
12

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
R
d-dimensional Euclidean space
gd
d-dimensional hypersphere
1l
Euclidean norm
)
standard dot product
gd
(d - diddmnnsiinnal hypersphere
0
unit vector
disjoint union
(X(
space of Lebesgue integrable functions on X
P(X)
space of probability distributions (or measures) on X
measures
8(-)
1-dimensional Dirac delta function
u($
uniform distribution Sd-1
pushforward (measure)
Y(
space continuous maps from to metric in metric space translation group of order d orthogonal group order d Euclidean group of order d element of group p- Wasserstein distance Sliced p- Wasserstein distance system of lines, tree system spherical ray spherical tree space symtems of & lines in IR
d(s-)
T(d)
((d)
E(d)
Wp
SW,
ry
T, T
yi....ll
L
L.
number of tree systems
k
number of lines in a system of lines or a tree system
R
original Radon Transform
Ra
Radon Transform Systems of Lines, or Radon Transform on Spherical Trees
(I: - d)-dimmnsio standard simplex
splitting map
Dirac delta function
space of tree systems
distribution space of tree systems
13

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
Supplemental Material for "Tree-Sliced Wasserstein Distance with Nonlinear Projection"
The supplementary is organized into four parts as follows:
In Section A, we provide background for Tree-Sliced Wasserstein distance.
In Section B, derive theoretical proofs for Radon transform on systems of lines with nonlinear projection.
In Section C, we describe Radon transform on systems lines for spherical functions.
In Section D, we provide further details for the experiments.
A. Background for Tree-Sliced Wasserstein Distance
In this section, we briefly outline the notion of the Radon Transform Systems of Lines (Tran al., 2025c) with its aetaace-basele extension (Tran al., 2025b).
Building blocks of Tree-sliced Wasserstein distance on Systems of I Lines. The Tree-sliced Wasserstein distance on Systems of Lines is constructed step-by-step as follows:
1. Given a positive number d presenting the dimension.
2. A line in Rd is element 1 (x, f) Rd x Sd-1_ Here, x is called the source and 0 is called the direction of the line.
3. A system of k lines in Idd an element of (Rd x gd-1yk. Denote a system of lines as C, and the space of all systems of
A lines by Ltd
4. point x in C can be parameterized as æ, 0g, where is the index of the line, and / the coordinate of the point on that zh lines.
5. A system of line L with additional tree structure is called a tree system (see (Tran et al., 2025c:b)). Each tree system is
a measure space, endowed with a tree metric.
6. A space of trees (collections of all tree systems with the same tree structure) is denoted by T with probability distribution a on T,, which comes from the tree sampling process.
7. For L E g. the space of integrable functions on L is:
& LT(C) = ( f: C C- R: lifll Σ
[f(z; +- < J
(23)
8. A splitting map is continuous map from Rd x Lg to the (k- 1)- dimensional standard simplex A i.e.
E C(R. x AL For E L.² (IRd), we define:
REf
C
R
(24)
+ 0 - / f(y) aly, 8 (y - 1y-x,,01 0,) dy.
(25)
The function R8f (C).
9. The operator:
Ra (RR))
II ('(C)
CET (REf)
f
is called the Radon Transform on Systems of Lines.
14

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
10. When the splitting map is E(d)-invariant (this B(d)-invariance will be described in the part), the Radon Transform Systems of Lines is injective (see (Tran et al., 2025b)).
11. The Tree-Sliced Wasserstein Distance on Systems of Lines, denoted by TSW-SL as in (Tran et al., 2025c), or its Distance-based variant Db-TSW as in (Tran et al., 2025b), between μ4, v in P(R4) is defined by:
Db-TSW(H, = / Wa R2fx) dC(
(26)
We choose the notation Db-TSW since this variant is generalization of TSW-SL. DD-TSW is identical with the definition of SW when & = 1, i.c. tree systems in Lg have only one line.
12. The Db-TSW distance is metric on P(R") (see (Tran et al., 2025b)).
13. It is worth noting that, on tree systems, optimal transport problems admits closed-form expression. since is a metric space with tree metric (see (Le et al., 2019)). Leveraging this closed-form expression and the Monte Carlo method, the distance in Eq. (26) can be efficiently approximated by a closed-form expression. Additionally, for the p-order Wasserstein with p 1, may consider the scalable variant-Sobolev transport (Le et al., 2022), which also yields closed-form expression for fast computation, and generalizes tree- tree- Wasserstein (i.e., 1-order Wasserstein tree) to more general settings such as for p > 1, and for measures on a graph.
The group E(d) and its action. The Euclidean group E(d) is the group of all transformations of Rd that preserve the Euclidean distance between any two points. the semidirect product between T(d) and O(d), i.e.
E(d) T(d) x O(d).
(27)
where T(d) is group of all translations in Rd and O(d) is the orthogonal group of Rd. Each element g of T(d) can be presented as pair:
g (Q,a) T(d) where Q O(d) and Rd.
(28)
The group E(d) acts on I2 naturally as follows: For x E Rd and = (Q,a) T(d), we have:
(g.z) - gz +
(29)
naturally induces a group action on the set of all lines in Rd, i.c. Rd x gd-1, For / = (xx 0) E Rd x Sd-1 and g = (Q)a) F(d), we have:
()
x ( S
(30)
For C {4: (z,, (z, 0,) E Lg. the action of E(d) on L is defined as:
gC = {g6 = (Q x, + a, Q 0,) E Lgg
(31)
The tree structure of tree system is preserved under the action of E(d) (see (Tran et al., 2025c;b)). In other words, if L e T a tree system, then gL is also tree system. The group action of E(d) on Lg induces a group action of F.(d) on T.
E(d) ((d)-iivariant splitting maps. splitting map a C(Rd Lg A 1) is E(d)-invariant,
algv = aly,
(32)
for all (y, L) E Rd LŞ and g E B(d).
Remark A.l. Equivariance is widely used in machine learning across various contexts, including equivariant models (Tran et al., 2024c) and equivariant metanetworks (Vo et al., 2024; Tran et al., 2024d;b;d). These approaches leverage symmetries in data or model architectures to improve generalization, reduce sample complexity, and ensure consistency under group transformations.
15

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
B. Theoretical Proofs for Radon Transform on Systems of Lines with Nonlinear Projection
B.1. CR is well-defined
We show that the Circular Radon Transform on Systems of Lines is well-defined.
Proof. Recall from Eq. (9), have:
CR (A(()) II (C) where / (CR2, CEL4.220*
(33)
r00
and CR f + (.) / f(v) aly, aly. (t lly =2 -r0:ll2) dy.
(34)
Rd
We have:
k |CCR2, fllc Σ
/
CR
f(2; +t. +t- dt
R
k
f(v) aly. L); . 5 lly - x; - -r0bll2) dy
( (Je ((() alv, L); 0 (t- - I1y - - rr0ill2) du) dt
Σ / (f ([ilw)) aly, C)s 8t Ily Is - -r0sll2) dt dy
Σ / [f(v)) a)y, (f (t - 11 - x, - r0,112) 0.-r0,112) dt dy
Σ { [f(v) a(v, dy
1=1 JRd
f 1f()) Ealu.0) dy
[ [f(v)] dy
11/
(35)
So CR E (C). implies that the operator CR2,; L²(RR) L-(C) is well-defined, as well as CR
Remark B. Note that, from the above proof, we see that f E )( i.e. f E L (RR)), Ilf = and f(y) > 0 for all
y E R.d. we also have ||CRE,, and CR2 + 0.) > for all m, + 02 L. implies that CR E P(C).
B.2. H is well-defined
We show that the Spatial Radon Transform on Systems of Lines well-defined.
Proof. Recall from Eq. (11), have:
:º: (Rd)
II ('(C) where f (Hzf) ceL
and Hzf((z + 1- 0.) = / f(y) a(h(v), L)x 8 (t (h(y) - I,, 0.) dy.
(36)
16

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
We have:
H
Σ
[122((; + 0.)..l dl,
1=1 JR
/Alee ((0) ((0) aab())) L) (1 (h(y) Fss 0,)) dy dl
Σ
(Le [f()) a(h(v), G); 5 (h(y) 21, 01)) dy dt
(f (116)) aa C); ( - (h(y) - nx 0.)) dt dy
/R"
Σ / a(k)),)))();))))) (f 5 (b(w) xz, 02) di dy
Σ JRd
[f(y)]
aah(y),))))) dy
i=1
k
S
If(w))- Σ a(h(v)), 1), dy
Jzd (f(w))
1lfll2
(37)
So H2. E (C). It implies the operator H2: L (RR)) (C) is well-defined, well 7
Remark B.2. Note that, from the above proof, that if E P(RR), i.e. f T.* (TR)), ||[III111111-- and f > 0 for all
y E Rd. also have ||H l and H2f(x, + t- ..) > 0 for all + 0 E C. It implies that H2f E P(C).
B.3. Proof for Theorem 4.2
Recall the original Circular Radon Transform CR (Kuchment, 2006; Kolouri al., 2019) follows:
CR
L (RR)
- L' (I x sd-1 x R20)
CRf,
(38)
where:
CRf:
R x Sd-1 x Ryo
R
(39)
J f(v) 5 (t 11 - -10l2) dy-
(t, 0,
(40)
In (Kuchment, 2006), it is showed that the Circular Radon Transform CR is injective. We will leverage this result to prove the injectivity of the proposed Circular Radon Transform on Systems of Lines CR
First, for each 0 E Sd-1 consider the tree system C consists of l: identical lines (0, Define the function y as follows:
g
R sd-1 x R200
(41)
(t,
CR C f(r + (200))))
(42)
Since
+ +t. = [J f(v) aly, L): 8 ( 11 - x; - -r0(1l2) dy.
(43)
17

---

Tree-Sliced Wasserstein Distance with Nonlinear Projection
We have:
gtt, (,,
Σ
CR
c6n _ + + 22m)))
(44)
{ f(v) aly, )()); 6t - 11 - ccnns -0 r00c)))ll2) dy
(45)
RRd
{ f(v) aly, [ ( - lly - r1112) dy
(46)
Jdd
Σ [ f(y) (t 11--y -lly -r0112) [ a(v. ((o))))))))) dy
k
(47)
RRd
J [(v) -8 5 1ly 10r01112) dy
Σ Rd
(48)
= CRf.
(49)
is clear that CR is a linear operator. To prove CR injective, consider / E Ker(CRD). By the definition of g in Eq. (41), we have g is the function 0. But g is exactly is the Circular Radon Transform of f. and since the Circular Radon Transform injective, we conclude that f the function O. In conclusion, CR injective.
B.4. Proof for Theorem 4.3
We present the proof for Theorem 4.3.
Proof. Recall the Radon Transform on Systems of Lines R (Tran et al., 2025b) as follows:
Ra L (RR)
II (C)
DELE
f (R2)
(50)
where
Ref
C
R
J f(v) aly, -):-((
D3 + t. 0:
8( - (y Ij, 0,)) dy.
(51)
It is proved in (Tran et al., 2025b) that RD is injective for B(d)-invariant splitting map We leverage this result to prove the Spatial Radon Transform on Systems of Lines is injective. First, by the injective continuous map h Id - Rd we show that the push-forward of f e IR via h. defined as:
hy)(1) { f(h (1h))()))
for all y E Rdo such that y E h(RR), for all y R such that y & h(Rd).
(52)
has its Radon Transform on Systems of Lines, i.e. (ba/h cEL equal to the Spatial Radon Transform on Systems of Lines of f, i.e. (H2S) cengo- In other words, for all L E L have:
H = R2
(53)
Indeed, we have:
(2) +t-0,) 0,) S/s hyf(v) aly, C)- 8 (t (y 2,, 0.)) dy
18

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
Sncee haf(y) oly, C) (t (y Tz, B()) dy
1 f(h f(h a(h())) G)r 8 (hly) - 24, 0()) dy
[ ((0) abbly) C) . (M(v) xz.0:)) dy
H2f(x; tt 0:).
(54)
is clear that 71"h is linear operator. To prove 7 injective, consider f E Ker(H("). Since H25 = R2 (2gf). implies that hyf E Ke Since R is injective, implies that hyf is the function 0. By the definition of the push-forward haf as in Eq. (52), we conclude that f is the function O. In conclusion, H is injective.
B.5. Proof Theorem 5.3
show that WircularTSW metric on P(Rd). The proof for SpatialTSW is similar.
Proof. We will show that:
CircularSS((p.2)
Wy. (C(),,
(55)
a metric on P(Rd), by verifying its positive definiteness, symmetry and triangle inequality. Positive definiteness.
For 14, P(Rd), is clear that
CircularTSW(ppp = O,
(56)
and
CircularTSW (m, v) > 0.
(57)
If CirculaTTS((x,b) = 0, then W(CR ff) = 0 for almost every L E Ty. Since W is a metric on P(C), we have CR = C for almost every C E T. By Theorem 4.2, implies that μ v.
Symmetry.
For M, E P(R"), we have:
(Ir NI. do (C)
Jg
W(CR CR2 (C)
T
CSrcularSSWpp...
(58)
So CireularTSW(u. v) CireularTSW(y.))
Triangle inequality.
For (11, 12, H3 P(R.D), have:
CireularTSW(H,, (MI, (2) + CireularTSW (H2, (3))
W(CR CRO (L) + / CRe (C(L)
Tt
[
W(CR MR CR2, (R de (L)
(C)
T
CireularTSW P3).
(59)
The triangle inequality holds for CircularTSW.
In conclusion, CircularTSW metric on the space P(R..
19

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
C. Radon Transform on Systems of Lines for Spherical Functions
In this section, we review (Tran al., 2025a) which proposes Spherical Radon Transform on Spherical Trees and Spherical Tree-Sliced Wasserstein distance, which are analogs to Radon Transform on Systems of Lines (Tran al., 2025c;b) and corresponding metric, applied for spherical functions. Then, we explain how to apply Generalized- framework in this paper for spherical settings
C.1. Background for Spherical Radon Transform on Spherical Trees
To make this easy to follow, we will follow the construction of Appendix A. Building blocks of Spherical Tree-Sliced Wasserstein distance.
1. Given a positive number d presenting the dimension. We will work with functions on the d-dimensional hypersphere gd [d+1, where:
gd == {x (1,..., E RR++1 : || = 1} c 12+11
Note that Sd is metric space with the metric dga defined as dga (a, b) arccos (a, p
2. The stereographic projection corresponding to x E gd is defined by:
Px
Sd {z}
Hg
(z, v) (x, v)
y
x
1 (x, y) y.
(60)
By convention, let (x:( = oo, then fax: Sd - Hz U {00}.
3. The spherical ray with root x and direction denoted by rig is defined as
+ - to (+l y: 1 (co))
(61)
Each ray q is isomorphic [0, x| via dga (x. +), so it is parameterized as (t, r.)..
4. Spherical trees T wy...... in gd is the gluing space of k spherical rays ri. at the root I. is the root and gi...., Uk are the edges of " It is measure metric space, endowed with tree metric.
5. The space of spherical trees with & edges in Sd is denoted by Ty. with probability distribution on Tg which comes from the tree sampling process.
6. For T TT the space of integrable functions T is:
* E (.( j: T -:T-R. R I1l Σ 1=1
ff 15(1, ((( < {o) (1l.)))
(62)
7. A splitting map is a continuous map from sd x Té to the (k 1)- dimennsioall standard simplex A i.c. a E C (Sd x Td. 1+-1). For f E L (8d), we define:
RGS
T
R
(63)
I f(y) aly, T)+ o(t arccos (z. v\) dy.
(t,
(64)
The function R$ is in (T).
8. The operator:
Rº: L'((") II L'(T) TEN f (R
is called the Spherical Radon Transform on Spherical Trees.
20

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
9. When the splitting map a is O(d + 1)iinvariant (this O(d + ))-ivaariance will be described in the next part), the Spherical Radon Transform Spherical Trees is injective (see (Tran et al., 2025a))-
10. The Spherical Tree-Sliced Wasserstein Distance (Tran al., 2025a) between ,4, P(Sd) is defined by:
ST) = [ a RFSu) (T(
(65)
11. The STSW distance is metric on P(S
12. It worth noting that, on tree systems, optimal transport problems admits closed-form expression, since is metric space with metric (see (Le et al., 2019)). Leveraging this closed-form expression and the Monte Carlo method, the distance in Eq. (65) can be efficiently approximated by a closed-form expression.
The group O + 1) and its actions. The orthogonal group O(d ((d+1) is the group of linear transformations of RR+11 that preserves the Euclidean norm ||. |2- The group O(d + 1) acts on @d naturally follows: For S and g Q O(d 1), we have:
(g, x) - gz Q x.
(66)
It naturally induces group action the set of all spherical lines in St, as well as spherical trees. The tree structure of spherical tree is preserved under the action of O(d + 1) (see (Tran et al., 2025c;a)). In other words, if T E T is spherical tree, then g is also a spherical tree.
Definition C.I. A splitting map a in C(Sd x TV 1 is said to be O(d + 1) invvarian,,, we have
algv. = aly, T)
(67)
for all (y, T) E gd To and E O(d 1).
A candidate for O(d + invariiat splitting maps is presented as follows: Consider the map B: gd x Tr - R
0,
if y - x or y
Bly, T
(u. v1)
(68)
V (2, if + tx.
arccos
(x,, w)²
The map 3 is continuous and O(d + invariian Take a: gd x Ty. - A to be:
aly, T) = softmax [3(y) T) D.)
(69)
C.2. Spatial Spherical Radon Transform on Spherical Trees
Consider positive integer dg, and injective continuous map h: Sd - Sdo and splitting map a C(T L AR-1)
defining the splitting mechanism. Let 'I be a spherical tree of k edges in Tye For function f E T.-( define the function H L (T) as follows:
R$S
T
IR
(70)
[ f(y) a(h(y); T): ö(t arccos (æ, h(())) dy,
(71)
The Spatial Spherical Radon Transform on Spherical Trees defined as the operator:
H
L (RR)
II L'(T)
TET
f
(H#
(72)
21

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
C.3. Proof for Theorem 4.4
We present the proof for Theorem 4.4 about the injectivity of the Spatial Spherical Radon Transform on Spherical Trees.
Proof. Recall the Radon Transform on Spherical Tres RD (Tran et al., 2025a) as follows:
R
L'(S")
II
('(T)
TETS
- (RTNTETT
(73)
where
R48
IR
J f(y) ((y,L), - 6 arccos (x, y)) dy,
(t, (.)))
(74)
is proved in (Tran al., 2025a) that R" injective for Old + i)-invariant splitting map a. We this result prove the Spatial Radon Transform on Spherical Trees is injective. First, by the injective continuous map h S gdo we show that the push-forward of Rd via h, defined as:
(v)) h4f(y) { 5(h (g))
{
for all y E Sde such that y E h(S*), for all E Sde such that $ h((d).
(75)
has its Spherical Radon Transform on Spherical Trees, i.c. {R7 (hg Terde equal to the Spatial Radon Transform on other words, for all T T we have:
Spherical Trees of f. f.c. (H f Terjo- Ter.
HSS R4(14))
(76)
Indeed, we have:
(t(h))((t l. hyf(y) aly, L), 5 6t arccos (z. y)) dy
-Kan hyf( aly, C): 5 arccos (z. v\) dy
[ f(h-1 a(h(y), C) 8 arccos (x, h(y))) dy
[ f(v)- -a((()(())))))))))) (): 5 arecos (z, h(y))) dy
H99 (t.
(77)
is clear that 7( is linear operator. To prove 70 is injective, consider f E Ker(H"). Since H -Rg it implies that hyf E Ker(R*). Since Ra is injective, implies that hy/ is the function 0. By the definition of the push-forward hy as in Eq. (75), conclude that f is the function O. In conclusion, H is injective.
C.4. Spatial Spherical Tree-Sliced Wasserstein Distance
For two probability measures y and with density function f/ and fu. Given positive integer de and choice of the continuous injective map h: odd - Sde the Spatial Spherical Tree-Slicee Wasserstein Distance between μ and is defined as the average Wasserstein distance the tree-metric space L between the Spatial Spherical Radon Transform on Spherical Trees of fµ and fu. Following (Tran et al., 2025a), this averaging is taken over the space of trees Té. according to distribution Thd which arises from the tree sampling process.
Definition C.2. The Spatial Spherical Tree-Sliced Wasserstein Distance between μ and v in P(S is defined by:
(patialSTS(H.vv J W(H H.F do
(78)
22

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
SpatialSTSW is metric the space P(S4) of measures on $4.
Theorem C.3. SpatialSTSW a metric on the space P(
Proof. We will show that:
/
(4at (1,,
5, (v) do dT(
(79)
a metric on P(S)), by verifying its positive definiteness, symmetry and triangle inequality.
Positive definiteness.
For Mr E P(()), it is clear that
(p) O,
(80)
and
SpatialSTSW(1,.)) v) >
(81)
SpatialSTSW( 0, then W(H fx) 0 for almost every T Ty Since W is metric on P(T), we have fz 169 fo for almost every E T. By the injectivity of the Spatial Spherical Radon Transform Spherical Trees, it implies that
Symmetry.
For u, E P(8d), we have:
(patialSTS((H.))) f HEFE) dT
HHH do
(patialSTSW(M.1))
(82)
So SpatialSTSS v) (1)
Triangle inequality.
For (4,, 12,, M4 E P(()),, we have:
SpatialSTSW(HH (41, + SpatialSTSW(Hz
W(H? Hzg do((T) + f.g. W(H H9 (wg) do (T)
Jup
Jope (W(H7fmm H9 W(H75mm H9F do (T)
W(H? H75 (T(
Jf
CireularTSW(H, ua).
(83)
The triangle inequality holds for SpatialSTSW.
In conclusion, SpatialSTSW is metric on the space P(3d).
The choice of the injective map h. Note that, the map h gd - gde has to satisfy the injective condition. We construct h as follows. We construct h as follows. First, consider do = d + 1. We define continuous function: k(y) = शामक म. Ed Ui + +e), which maps y e gd to the range (0, .T. We set e = 10 6, Using this, we define the mapping h(y) (cos(k(y)), sin(k(y)) v), which is injective.
23

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
D. Experimental Details
D.I. Algorithm of proposed Tree-Sliced Distances
We describe the pseudo-codes for CireularTSW, SpatialTSW, SpatialSTSW in Algorithms 2, respectively.
Algorithm 1 Circular Tree-Sliced Wasserstein distance.
Input: Probability measures i and in P(R"), number of tree systems L, number of lines in tree system k space of tree systems T, splitting maps a, and parameter E Rgo.
for i = 1 to L do
Sampling x E Rd and .1..... 0g '2d 1(
Contruct tree system L; = ((1, 0.),..., (z)
Projecting μ and v onto L; get CR2 and 2
Compute (p,, - (L/L) end for Return: (IreularTSW(p,v))
2 CR 2).
Algorithm Spatial Tree-Sliced Wasserstein distance.
Input: Probability measures f and v in P(Rd)), number of tree systems L,, number of lines in tree system k, space of tree systems T, splitting maps and injective continuous map h IR - R
for i 1 to L do
Sampling x E Rd and 01..... 0g iid U(S
Contruct tree system L; ((1, 01) f1),..., (s,, Az))- Projecting μ and v onto T; to get H2, p and H2, v. Compute SpatialTSW(j,v)) (1/L) M, end for Return: SpatialTSW(y,.u (,,
Algorithm Spatial Spherical Tree-Sliced Wasserstein distance.
Input: Probability measures y and v in P(S)), number of tree systems 7., number of lines in tree system k, space of tree systems T. splitting maps a, and injective continuous map / : Sd - gdo
for 1 to L do
Sampling x E Sdle and yi,.... yk Lid (S
Contruct tree system T; {(x, ((x,)),... yi).... (s, y- Projecting μ and v onto L: get H and V. Compute SpatialSTSW(H..)) (1/L) W(H9,H, H4.y) end for Return: SpatialSTS((ppv).
D.2. Computational and Memory Complexity
We provide complexity and memory analysis of our proposed distance. Since memory on GPU can be optimized for parallel processing capabilities, we provide the empirical memory usage GPU, with expectation that our distance would be used in GPU setting which is standard in machine learning.
Computation and Memory Complexity. Assuming > m, the computational complexity of SpWtialTSW and CircularTSW (Lkndd + Lkn log n), while CircularTSW. and SpatialSTSW have a more efficient complexity of O(Lnndg + Ln logoogg n). All distances share an empirical memory cost O(Lkn + Lkdg + nde). We analyze the main operations of our proposed distances in Table 6.
24

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
Table 6: Computational and Memory Complexity Analysis of proposed Tree-Sliced distances
Distance
Operation
Description
Computation
Memory
Mapping
Map points onto new space Matrix multiplication points and lines
O(ndy))
O(nis)
Projection
O(Lkndo) O(Lmnds)
O[Lkds + nie)
SpatialTSW
Distance-based weight
Distance calculation and softmax
O + Lkde (id))
splitting
Sorting
Sorting projected coordinates
OLLm log O + Lkm log
O(Lkn)
Total
O Lkdg ndg)
Circular projection
Subtraction und Norm culculation Distance calculation and softmax
(LLndd)) (LLnud)
O + nide)
CisWulaTTSW
Distance based weight
O + Lkdda ndo)
splitting
Sorting
Sorting projected coordinates
O log O[Lkndo Lknlogg)
O(Lkn)
Total
O + Lkd, ndo)
Circular projection
Subtraction und Norm calculation Distance calculation and softmax
((Lnds) ((Lknd))
O(Lda nde)
CiruuarTTSW
Distance-huu -aaeed weight
0(1km Lkde nds)
splittnng
Sorting
Sorting projected coordinates
O(Lnlogn)
O(Ln)
Total
OLLnndo La logn)
O + Lkde ndo)
Mapping
Map points onto space Matrix multiplication of points and source Distance calculation softmax
((udy)
O(nds)
Projection
O Luds)
O + nde)
S
Distance-based weight
OLLAnds)
O + Lkdo ndo)
splitting
Sorting
Sorting projected coordinates
OL n. O (L.nnd log
O(Ln)
Total
O(Lkn + nds)
Projection and Sorting, In CircularTSW,-o and SpatialSTSW, the projected coordinates within a tree is the same for all lines. As result, the computational complexity of these two steps in CircularTSW, and SpatialSTSW is reduced by a factor of the number of lines in a tree, k, compared to SpatialTSW and CircularTSWW This reduction the primary reason for the computational advantage of CireularTSW,=00 and SpatialSTSW.
Memory Cost of Distance-Based Splitting. As previously noted in prior work (Tran al., 2025b), the empirical GPU- optimized memory cost of distance-based splitting is lower than its theoretical estimate due to kernel fusion optimizations. This operation consists of: (1) computing distance vectors from points to lines (O(Lkndg) computation and memory), (2) calculating their norms (O(kn)da) computation and O(Lknda) memory), and (3) applying softmax over all lines in each tree (O (Lkn) computation and memory). While the thcoretical cost O(Lkndg) for both computation and memory, we leverage PyTorch's automatic kernel fusion (via 'torch.compile') to merge these steps into single operation. This enables the distance vectors (Lkn dg) be stored in shared GPU memory rather than global memory. As result, only three matrices need be stored: line matrix (O(Tkda))), support matrix (O(ndo)), and split weight matrix (O(Tkn)). reducing overall GPU memory usage OLLkn + Lkda + nda).
D.3. Runtime and memory analysis
d= 10
20000
d=
50
17500
50
Memory Us3ge (MB)
d= 500
d= 500
Runtime (ms)
d= 1000
12500
1000
10000
7500
5000
2500
10000
20000
50000
10000
40000
50000
Number of supports (n)
Number of supports (n)
Figure 3: Runtime and memory evolution of SpatialTSW.
In this section, we conduct runtime and memory analysis of SpatialTSW, CircularTSW. and CircularTSW. , with
25

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
500
d= 10
4000
d= 10
50
400
Memory usage (MB)
Runtime (ms)
d= 500
500
1000
1000
2000
100
1000
10000
20000
30000
50000
10000
40000
50000
Number of supports (n)
Number of supports (n)
Figure 4: Runtime and memory analysis of Circular
1400
10
d
1.4
d= 50
Merrory usage (MB)
Runtime (ms)
1.2
= 500
1000
d= 500
1000
1000
800
1.0
0.8
0.6
200
10000
20000
30000
40000
50000
10000
20000
30000
40000
50000
Number of supports (n)
Number of supports (n)
Figure 5: Runtime and memory analysis of CireularTSWW -=
respect to the number of supports and the support's dimension on a single NVIDIA A100 GPU. We fix L = 2500 and k 4 (following the practical setting from our diffusion model experiments) for all configurations and vary E {500, 1000, 5000, 10000, 50000} and E {10, 50, 100, 500, 1000].
Runtime scaling. Figures and demonstrate linear relationship between the runtime of SpatialTSW and CircularTSW and the number of supports Regarding scaling with the data dimension, we observe that d = 10000 takes approximately twice the runtime of d 5000, suggesting linear relationship between dimension and computational time. This linear trend aligns with our theoretical complexity analysis in Appendix D.2. It is also worth noting that CircularTSW runs faster than SSW, as it relies on vector norms instead of vector multiplications. Regarding CircularTSW, Figure 5 also demonstrates an almost linear relationship with the number of supports n. Interestingly, the runtime of SircularTSSW,... scales very efficiently with d. We suspect this due to the reduced amount of vector normalization required compared to CircularTSW (by a factor of k). Ultimately, CircularTSW.-. is significantly faster than other Tree-Sliced methods, by two orders of magnitude when d and are sufficiently large, which aligns with our theoretical complexity analysis. The significant reduction in computational cost when using the Circular Sliced Wasserstein variants arises from the efficiency of computing L2 norms compared to inner products. This advantage is illustrated in Figure 6.
Memory scaling. Figures 4, and presents the memory consumption analysis of SpatialTSW, CircularTSW, and SrcullrTSWW,.a revealing linear scaling relationship with and n. This aligns with the theoretical complexity analysis and suggests a predictable scaling behavior. Notably, the peak memory usage CircularTSW and CireularTSW. , is significantly lower than that of SpatialTSW.
26

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
Vector Inner Product
1.0
Runtime
0.4
Number of Vectors
Figure 6: Rebuttal result comparing inner product and L2 norm performance, benchmarked on an A100 GPU with fixed vector dimension of d 3000. The 1.2 norm demonstrates significantly faster computation times as the number of vectors increases.
D.4. Splitting map for CireularTSW
We recall that the splitting map is defined as:
aly, C); softmax (d(y, )=)
(84)
where d(y, L), represents the distance between and the 7th line in L.
In the context of the Radon Transform on a system of lines, this involves computing the inner product between the support and the projection directions. However, in the Circular Radon Transform system of lines, the projection coordinate calculation involves the Euclidean norm
||2-
Therefore, we define more computationally efficient distance function as:
d(y, L); = lly x; ly r; r0:ll2 0allg
(85)
where x; is the source, is the direction of the 7th line in C. and a fixed radius.
Notably, this distance function still results in an invarrinvari splitting map.
D.5. Analysis number of lines in CircularTSW...
In this section, we demonstrate that CireularTSW-DD.aa effective only tree-based setting. To illustrate this, we conduct experiment using the Denoising Diffusion Generative Adversarial Network (DDGAN), following the setup described in Appendix D.6. Figure 7 presents the FID scores over 300 training epochs for models using CireularTSW,-n. with A: = 1 and & = 4. The results clearly show that only the model with k = 4 trains stably, with its FID score gradually improving over time. In contrast, for = 1, the FID score suddenly spikes to 500 after epoch 125, indicating that the model collapses and starts generating meaningless images (all black pixels). This confirms that CircularTSW, is specifically designed for tree-like structures, relying on distance-basee based splitting map to function effectively.
D.6. Denoising Diffusion Generative Adversarial Network
Diffusion Models. Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020) have gained popularity as powerful gen- erative models capable of producing high-quality data. this experiment, we introduce their mechanisms and demonstrate
27

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
400
FID Score
300
200
100
CicularTSW with K-4
Creular with =
100
150
200
250
300
Epochs
Figure FID scores over the training process for CircularTSW-DD. with and
the improvements introduced by our approach. The diffusion process begins with initial sample from the distribution ((2)) and progressively corrupts it by adding Gaussian noise over T steps. This process is formally defined as:
T
(<11r(=)) = Il a(()((=11)
where each transition follows a Gaussian distribution:
(( N(ri v1 1,1),,
with predefined variance schedule Se.
The objective of denoising diffusion models to learn the reverse diffusion process, which reconstructs the original data from noisy samples. This requires estimating the parameters 0 of the reverse process, formulated
po(oor) = p(er) II po po(
where each step follows Gaussian transition:
Training these models typically done by maximizing the evidence lower bound (ELBO), which minimizes the Kullback- Leibler (KL) divergence between the true posterior and the model's approximation the reverse diffusion process. This expressed as:
Eg())) ( (po (pr + C,
where KL(-II-) represents the Kullback-Leibles divergence, and C' is constant term.
28

---

Tree-Sliced Wasserstein Distance with Nonlinear Projection
Denoising Diffusion GANs. While diffusion models generate high-quality and diverse samples, their slow sampling process limits real-world applicability. Denoising Diffusion GANs (DDGANs) (Xiao et al., 2021) address this issue by modeling each demoising step using multimodal conditional GAN, allowing for larger denoising steps. This significantly reduces the number of steps to just 4, leading to sampling speeds up to 2000 times faster than traditional diffusion models while maintaining competitive sample quality and diversity. The implicit denoising model in DDGANS is formulated as:
pa pa( = / pe (Zt (zr ( e) e)de, e~,10,
Xiao et al. (2021) optimize the model parameters 0 using adversarial training, with the objective:
T
min Σ (a
where Dadv represents the adversarial loss. Instead, Nguyen al. (2024b) replace the adversarial loss with the Augmented Generalized Mini-batch Energy (AGME) distance. For two distributions μ and given a mini-batch size > 1, AGME using a Sliced Wasserstein (SW) kemel is defined as:
AGME, (U, vig) = GME B),
where ji fal and i fiV. with f(x) = (Ir, (()) for nonlinear function g R - -R. The Generalized Mini-batch Energy (GME) distance (Salimans et al., 2018) is defined as:
GME = Btt.::--BBDD.../----D.----------
i...d where X, X ~ 1m and Y, y iidd yom with
Px
X = (m))
Here, D represents any valid distance metric. In our work, we replace D with Tree-Sliced Wasserstein (TSW) variants (our methods) and Sliced Wasserstein (SW) variants.
Setting. We adopt the same architecture and hyperparameters as Nguyen et al. (2024b) and Tran et al. (2025b). Our models are trained for 1800 epochs. For Tree-Sliced methods, including our own, we L 2500 and k = 4. For vanilla SW and SW variants, we follow Nguyen et al. (2024b) and use L 10000. The learning rate is also set according Nguyen et al. (2024b), where Ir.- = 1. 1.25e-4 and Ir g = 1.6c-4. For SpatialTSSW, we define hyy) + y², and for CircularTSW. we set r = 0.01. The standard deviation in tree sampling follows Tran et al. (2025b) and is set to 0.1. To evaluate runtime, we use a batch size of 64 and measure time on a single NVIDIA A100 GPU.
Qualitative Results. Figure 8 presents the qualitative results of SpatialTSW-DD, CireularTSW-DD. and Cireular,=o TSW: DD.
D.7. Gradient Flow
Detailed results on the 25 Gaussians dataset. Table presents the detailed results of our proposed methods and baselines. The low standard deviation of SpatialTSW indicates that our method consistently achieves faster convergence compared other methods.
Ablation on h Rd - ds in SpatialTSW. We ablate several injective continuous functions h in SpatialTSW on the
25- Gaussian dataset. The results in Table show that, in general, h(y) = + yy² outperforms h(y) = y 77 although
the latter tends to converge faster during the first 1500 iterations. The best result is achieved with h(y) = y + 0.5g3, yielding final W2 value of 9.
29

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
SpatialTSW-DD
CircularTSW-DD
Circular, TSW-DD
Figure 8: Example images generated by our proposed DDGAN. The images correspond to (Left) SpatiaITSW-DD, (Middle) CircularTSW-DD, and (Right) orocula,=oTTWWDD..
Table 7: Detailed results the 25 Gaussians dataset. The table reports the average Wasserstein distance between source and target distributions over runs.
lteration
Methods
Time/Iter(s)
500
1000
1500
2000
2500
SW
4.21c-1 5.39c-3
1.54c-1 + 2.43e-3
7.72c-2 3,88e-3
4.97c-2 + 3.30c-3
3.59c-2 3,43e-3
0.0018
MaxSW
5,23e-1
8.31e-3
2.36e- 4.63e-
1,23e-1 3.17e-
8.04e-2 3.70e-
6.76e-2 3.07e-
SWGG
6.59e-1 1.93e-2
3.62c-1 4: 2,70e-2
1.92e-1 1.99-22 2.26e-2 1.39e-3
9.07e- 1.31e-2
4.42e-2 1.90e-2
0.0019
LCVSW
3.46e-1 4.63e-3
6.96e:2 + 3.11e=3
1.31e: +: 2.07e=3
9.28e-3 9.25e-4
0.0019
TSW-SL
3.49e-1 4.61e-3
8.10e-2 1. 2.34e-3
1.06e-2
1.00e-3
2.68e:3 1. 3.24e-4
3.16e-6
1.99e-6
Db-TSW
3,50e-1 5.10e-3
8.12e-2 1 2.34e-3
1,09e-2 1.41e-3
1.77e-3 1: 6.69e-4
1.30e-7
9.28e-9
0.0020
Db-TSW-
3.52e-1
5.17e-3
7.69e:2 + 3.37e-3
2.73e-2 4.87e-4
2.56e:3 + 7.72e-4
2.03e-6
3.70a-6
0.0021
SpatialTSW
3,20e-1 1 4.73e-3
3.44e-2 1 2.42e-3
2,95e-3
1.46e-4
3,97e-4 1 8.13e-5
1.17e-7
2.24e-8
0.0021
CireularTSW
4,28c-1 4.33e-3
1.20e-1 + 2.37e-3
3.48e-2 5.10044
1.41e-2 + 7.50c-4
7.86c-3 3.94e-4
0.0017
CircularTSW,....
4.32c-1
4.01e-3
1,226-1 + 1.50e-3
3.4le-2 2:226-3
1.45e-2 + 1.03e-3
8.94e-3 9.42c-4
0.0015
Hyperparameters. For Tree-Sliced methods, we set L = 25 and 4. For the SW and SW-variant baselines, we use = 100. The global learning rate is set to 0.001. Each distribution in both datasets sampled 500 supports.
D.S. A Guide to Selecting Projection Variants
Our motivation for proposing the non-linear projection framework is inspired by Generalized Sliced- Wasserstein (GSW) (Kolouri et al., 2019), which also includes both Circular and Spatial variants. underexplored in prior studies that among the three cersions-origtn SW, SpatialSW, and CireularSW, which variant is most suitable for a given task.
This suggests that among the corresponding TSW variants, such as Db-TSW (Tran al., 2025b), CircularTSW: and Spatial TSW, there is no guarantee that the versions with non-linear projections will consistently outperform the linear-props iinjer-rooecction TSW. However, the two new distance variants each offer distinct advantages over standard TSW, as outlined below:
The definition of SpatialTSW subsumes Db-TSW as special case when the function is the identity map. This implies that models leveraging SpatialTSW have, in theory, greater representational capacity than those using Db-TSW. A similar relationship holds between the corresponding SW variants.
The definition of CircularTSW is theoretically non-comparable to Db-TSW due to their fundamentally different constructions. However, CircularTSW. offers improved runtime efficiency. This benefit does not hold in the Sw context, where CircularSW, performs poorly. One reason is that CircularSW defines only a pseudo-metie metric, while Cireular r is true metric.
30

---

Tree-Sliced Wasserstein Distance with Nonlinear Projection
Table 8: Ablation study on the choice of h in SpatialTSW for Gradient Flow. Results show the average Wasserstein distance between source and target distributions over 5 runs on the 25 Gaussians dataset.
Iteration
k(y)
?
500
1000
1500
2000
2500
0.1
3.49c-1
7.12c-2
7.77c-3
6.49c-4
1.15c-7
0.5
3.33e-1
4.68c-2
3.74e-3
1.58e-7
9.59c-8
4+ 79"
1
3.20e-1
3.44--2
2.95e-3
3.97e-4
1.17e-7
5
2.88c-1
3.26e-2
3,57e-3
3.39e-4
2.24c-7
10
2.73e-1
3,35e-2
5.11e-3
1.54c-4
5,59c-7
0.1
3.57e-1
7.29e-2
7.72e-3
1.71e-3
1.32e-7
0.5
3.36e-1 1
4.24e-2
4.57e-3
1.06e-3
1.50e-7
4 79"
1 2.99e-1
2.37e-2
2.95e-3
3.50e-5
1.63e-7
5
1.75e-1
8.54e-3
2.60e-3
2.03e-3
1.85e-3
10
1.34e-1
6.30e-3
3.14e-4
6.40e-6
1.55e-6
Our framework offers greater flexibility by enabling broader selection of distance functions. However, in Machine Learning, predicting the best variant for a task often requires empirical experimentation. Table 9 shows that both Db-TSW and SSW perform well, but the non-linearity in SSW makes it hard to determine in advance which variant is better suited for a given task.
We offer intuition for selecting CircularTSW and Circular Since these distances rely on the L2 norm for the projection step, they are likely perform well when the L2 norms of the data are diversely distributed. We validate this advantage over Db-TSW and SpatialTSW in the Table 10, where the distribution of L2 norms is uniform. We speculate that this property explains why CireularTSW performs effectively for the Diffusion experiment (Table 1).
To the best of our knowledge, Db-TSW (Tran et al., 2025b) is the only tree-sliced distance effectively suited for large-scale generative tasks involving transport from a training measure to a target measure in Euclidean space. Previously, (Tran et al., 2025c) presents basic and limited version of (Tran et al., 2025b), primarily emphasizing the constructive aspects of the tree-sliced approach, which serve as foundational groundwork. Meanwhile, (Tran et al.. 2025a) explores the method in a spherical setting. Other works Tree-Sliced Wasserstein (TSW), such as (Le al., 2019), (Yamada et al., 2022), and others, are mainly designed for classification tasks and are not applicable to generative settings. This limitation arises because these methods rely on clustering-basee framework for computing slices, which is theoretically unsuitable (as the clustering must be recomputed each time the training measure is updated, rendering previous clustering results irrelevant) and empirically inefficient (since clustering is significantly more computationally expensive than linear or non-linear projection methods).
Table 9: Results for Tree-Sliced variants (Linear, Spatial, Circular) in a Gradient Flow task across datasets, showing the average Wasserstein distance between source and target distributions over 5 runs. Each method uses 100 projecting directions, trained for 500 iterations, with the best result reported over Ir E (1, 5 10-1, x 10-1, 1 x 10-2, 5,5x x 10-2 10-3,3 3 3 x 10- 5 x 10 Sp performs best on Half Moons, Swiss Roll, 25 Gaussians, and 8 Gaussians datasets, while Db-TSW excels on the Circle dataset.
10
Methods
Circle
Half Moons
Swiss Roll
25 Gaussians
8 Gaussians
SW Db-TSW
6.463e-4 + 5.1
1.648c-4 + 3.754e-5
9.795c-4 1.328c-4
4.007e-2 + 1,940c-3
3.786e-2 + 7,090c-3
1.331e + 1.123e-7
1.659c-6 + 1.47lc-7
1.659c-6 1.47le-7
2.103e-3 6.378e-4
3,475e-3 1.151c-3
SpatialSW
2.098c-4 5
2.146c-4 + 5.308e
9.329c + 1.113e-4
5.206e-2 2.456e-3
3.593e-2 + 2.619e-3
SpatialTSW
1 9
1.267e-6 6.638e-8
1.615e-6
1.751e-7
1.969e-3 6.314e-4
2.652e-3 6.996e-4
CircularSW
6.044e 1.670e-5
8.147e-5 5: 4 6.982e-5 4.253e-6
7.950e + 1.065e-4
1.150e-1 4.502e-3
2.137e-1 1.276e-2
Cir
1.493e-5 1.922
4
1.172e-2 8.564e-4
1.307e-2 1.531e-3
StrcularTSW....
9.201e-5 1 9.562e-6
4
2.009e-2 1.612e-3
3.044e-2 9.105e-4
Hy pyprrrarrmeerr r. Selecting the optimal hyperparameter, such as r for CircularTSW. challenging and often requires empirical tuning. Intuitively, should be large enough to ensure diverse projections onto the lines but should not exceed the
31

---

Tree-Sliced Wasserstein Distance with Nonlinear Projection
Table 10: Results on the advantage of the Circular Tree-Sliced variant in a Gradient Flow task when the L2 norm of the data is uniformly distributed. Data is sampled such that the L2 norm follows Uniform(0. 1). The table reports the average Wasserstein distance between source and target distributions over 5 runs. Each method uses 100 projecting directions and is trained for 500 iterations, with the best result reported over Ir {1,5 10-1, x 10-!, x 10-2, 5 10-2, x 10-$, 3 x 10-3, 5 10-3). CircularTSW consistently achieves lower Wasserstein distance, while Linear and Spatial variants struggle as the dimension d increases.
Methods
2000
d = 5000
d 10000
SW
0.535 + 0.006
9.795 + 0.025
70.06 0.100
Db-TSW
4.871 0.049
87.49 + 0.137
308.97 + 0.367
SpatialSW
1.510 +: 0.006
18.66 + 0.043
95.55 + 0.170
SpatialTSW
6.394 0.066
93.08 1 0.131
314.44 + 0,269
CircularSW
0.357 0.005
0.404 1 0.007
0.428 1 0.004
Cwr
0.304 0.009
0.347 1 0.010
0.369 0.015
CircularTSW.....
0.332 = 0.010
0.517 = 0.015
0.873 = 0.022
data's magnitude. For normalized data, we suggest starting with va
1 and tuning from there.
D.9. Spherical Gradient Flow
Data. Given the probability density function of the von Mises-Fisher distribution [1; = Ca((s) exp(xp² z:), where
4/22-1
E Sd is mean direction and K > 0 is concentration parameter and the normalization constant Cd(k) = (2 (2 1 1/2-1
we use 12 vMFs as the target distribution with x = 50 and mean directions as follows:
f1 = 6,
#2 = (1, ó 0), μs = (0, 1,
F3 = -))
#1 = (1, (1,00.0)
Ms (0, Ø),
As (0, (0))
As (0. o)
14 (ø, 0, -1),
Pio (,)
M11 (-0,0, -1)), M12 (-0, 0,0,),,
where o = 1+ v5
2
Similar to Tran et al. (2024a; 2025a), we pick 200 samples from each vMF.
Setting. We set L = 200 trees and lines for STSW and SpatialSTSW, and L, 1000 projections for other sliced methods. ARI-S3W (30) employs 30 rotations with pool size of 1000, whereas RI-S3W (1) and RI-S3W (5) use 1 and 5 rotations, respectively. Training is conducted using Adam (Kinga et al., 2015) optimizer with learning rate Ir 0.01, following update rules (Bonet et al., 2022):
xk++1) = (k)) V (2) SpatialSTSW(i 1/),
x(++1) 1
(k++))
D.10. Self-Supervised Learning
Encoder. Following the approach outlined in Bonet et al. (2022); Tran et al. (2024a; 2025a), we use ResNet18 (He al., 2016) as the encoder. We train on CIFAR-10 data for 200 epochs with a batch size of 512 and the SGD optimizer with initial Ir 0.05, momentum 0.9, and weight decay 10-3 To generate positive pairs, we employ commonly used augmentation techniques, consistent with previous studies (Wang & Isola, 2020; Bonet et al., 2022; Tran et al., 2024a; 2025a). These transformations include resizing, cropping, horizontal flipping, color jittering, and random grayscale conversion.
For STSW and SpatialSTSW, we configure L = 200 trees and k = 20. For other distances, we use L = 200 projections and NR = with a pool size of 100 for RI-S3W and ARI-S3W. The regularization coefficients are chosen as follows: À = 10 for STSW and SpatialSTSW. À = for SW, À = 20 for SSW, and À = 0.5 for S3W variants.
32

---

Tree-Sliced Wasserstein Distance with Nonlincar Projection
Linear Classifier. To evaluate the quality of the feature representations learned by the pre-trained encoder, linear classifier is trained top of these features. Following the approach of Bonet et al. (2022), the classifier is trained for 100 epochs using the Adam (Kinga et al., 2015) optimizer. We set the initial learning rate to 10 'ogethhrr with weight decay of 0.2 at epochs 60 and 80.
D.11. Sliced-Wassersterr Wasserstein Autoencoder (SWAE)
Adam (Kinga al., 2015) optimize, setting learning rate to Ir 10-3, The model undergoes training over 100 epochs with a batch size of 500, where the binary cross-entropy (BCE) loss function serves as the reconstruction loss. For STSW and SpatialSTSW, we fix L = 200 trees and k = 10 lines. Other sliced methods use = 100 projections. RI-S3W ad ARI-S3W use NR = 5 rotation and pool size of 100. We use prior 10 vMFs, while setting the regularization parameter for STSW, SpatialSTSW,
Setup. For model training,
10 for SSW, and 10-3 for SW and S3W variants.
CIFAR-10 Model Architecture. (Tran al., 2024a; 2025a)
Encoder:
E 33x32x32
Conv2dg2 ReLU + Conv2d32 ReLU
Conv2de4 ReLU + Conv2de4 ReLU
Conv2d12s ReLU - Conv2d128 Flatten
FC512 - ReLU - FC3
l normalization - E g
Decoder:
E E
FC512 FC2048 - ReLU
Reshape(128 x 4 x 4) - Conv2dT - ReLU
Conv2dT61 - ReLU - Conv2dT64 - ReLU
Conv2dT32 ReLU - Conv2dT32 - ReLU
Conv2dT3s Sigmoid
D.12. Hardware settings
The gradient flow experiments were conducted on a single NVIDIA A100 GPU, with each experiment taking approximately 0.5 hours. The denoising diffuuionn experiments were executed in parallel NVIDIA A100 GPUs, with each run lasting around 50 hours. All spherical experiments were conducted single NVIDIA A100 GPU.
33


Reviews:
[{"review_id": "EVFWx5oEBh", "reviewer": "Reviewer_m8Ge", "summary": "The authors introduce several new variants of tree-sliced Wasserstein distance, which was introduced in [TPTLN '24]. This is done via two new proposed Radon transforms: (1) the generalized Radon transform on systems of lines and (2) the spatial Radon transform on systems of lines. Using their new Radon transforms, the authors define two variants tree-sliced Wasserstein distance: (1) circular tree-sliced Wasserstein (CircularTSW) distance and (2) spatial tree-sliced Wasserstein (SpatialTSW) distance. Unlike previous work, which used linear projections to construct tree-sliced Wasserstein distance, CircularTSW and SpatialTSW incorporate non-linear projections. \n\n\"Tree-sliced wasserstein distance on a system of lines\" [TPTLN '24]\n\n### Update after rebuttal\n\nThanks to the authors for their response. I maintain my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We direct the Reviewer to Table R1-2 at https://sites.google.com/view/nonlinear-tsw.\n\n**W1. [...] (many variants of TSW) outperform SW or TSW with linear projection**\n\n**Q1. [...] (intuition) one (TSW) over another**\n\n**Answer.** Our motivation for proposing the non-linear projection framework is inspired by Generalized Sliced-Wasserstein (GSW) [1], which also includes both Circular and Spatial variants. It is underexplored in prior studies that among the three versions—original SW, SpatialSW, and CircularSW,  which variant is most suitable for a given task. \n\nThis suggests that among the corresponding TSW variants—Db-TSW [2], CircularTSW, and Spatial TSW—there is no guarantee that the versions with non-linear projections will consistently outperform the linear-projection TSW. However, the two new distance variants each offer distinct advantages over standard TSW, as outlined below:\n\n- The definition of SpatialTSW subsumes Db-TSW as a special case when the function $h$ is the identity map (see lines 122–124). This implies that models leveraging SpatialTSW have, in theory, greater representational capacity than those using Db-TSW. A similar relationship holds between the corresponding SW variants.\n- The definition of CircularTSW is theoretically non-comparable to Db-TSW due to their fundamentally different constructions. However, $CircularTSW_{r=0}$ offers improved runtime efficiency. This benefit does not hold in the SW context, where $CircularSW_{r=0}$ performs poorly (see lines 252–261). One reason is that $CircularSW_{r=0}$ defines only a pseudo-metric, while $CircularTSW_{r=0}$ is a true metric.\n\nOur framework offers greater flexibility by enabling a broader selection of distance functions. However, in Machine Learning, predicting the best variant for a task often requires empirical experimentation. Table R1 shows that both Db-TSW and SpatialTSW perform well, but the non-linearity in SpatialTSW makes it hard to determine in advance which variant is better suited for a given task.\n\nWe offer intuition for selecting CircularTSW and $CircularTSW_{r=0}$. Since these distances rely on the $L_2$ norm for the projection step, they are likely to perform well when the $L_2$ norms of the data are diversely distributed. We validate this advantage over Db-TSW and SpatialTSW in Table R2, where the distribution of $L_2$ norms is uniform. We speculate that this property explains why CircularTSW performs effectively for the Diffusion experiment (Table 1).\n\nTo the best of our knowledge, Db-TSW [2] is the only tree-sliced distance effectively suited for large-scale generative tasks involving transport from a training measure to a target measure in Euclidean space. Previously, [3] presents a basic and limited version of [2], primarily emphasizing the constructive aspects of the tree-sliced approach, which serve as foundational groundwork. Meanwhile, [4] explores the method in a spherical setting. Other works on Tree-Sliced Wasserstein (TSW), such as [5], [7], and others, are mainly designed for classification tasks and are not applicable to generative settings. This limitation arises because these methods rely on a clustering-based framework for computing slices—which is theoretically unsuitable (as the clustering must be recomputed each time the training measure is updated, **rendering previous clustering results irrelevant**) and empirically inefficient (since clustering is significantly more computationally expensive than linear or non-linear projection methods).\n\n**Q2. [...] (select  $r$) affect performance empirically?**\n\n**Answer Q2.** Selecting the optimal hyperparameter, such as $r$ for CircularTSW, is challenging and often requires empirical tuning. Intuitively, $r$ should be large enough to ensure diverse projections onto the lines but should not exceed the data's magnitude. For normalized data, we suggest starting with $r = \\frac{1}{\\sqrt{d}}$ and tuning from there.\n\n**Q3. [...] (include) CRT and SRT?**\n\n**Answer Q3.** The TSW-SL [3] offers more general tree structure than the concurrent-lines tree structure used in [2]. However, this generality comes at the cost of runtime efficiency, as the concurrent-lines structure allows a GPU-friendly implementation. Since efficiency is a priority in our work, we adopt the concurrent-lines structure.\n\nWhen extending TSW-SL to non-linear projections, we initially believed SRT could apply to general tree structures, as key properties like injectivity are preserved. However, this does not seem to hold for CRT, since parts of the proof that CircularTSW defines a valid metric rely on the specific concurrent-lines structure.\n\n---\nWe thank the Reviewer for the constructive feedback and for pointing out the typos, which we will address. If our clarifications are satisfactory, we kindly ask you to consider raising the score. We are happy to address any further concerns in the next discussion stage.\n\n---\n*References.* Kindly refer to **References** in our response to Reviewer jNth.", "claims_and_evidence": "All claims/theorems are supported by proofs.", "methods_and_evaluation_criteria": "The benchmark datasets seem reasonable to me.", "theoretical_claims": "I did not check the proofs.", "experimental_designs_or_analyses": "I checked the experimental design and it seems reasonable to me.", "relation_to_broader_scientific_literature": "This paper is a follow up to [TPTLN '24] and [TCNLN '25]. While [TPTLN '24] introduces the general tree-sliced Wasserstein distance on systems of lines, this work extends it by introducing generalized Radon transforms on systems of lines and defining the corresponding tree-sliced Wasserstein distances.  The current paper extends previous work by integrating nonlinear projection mechanisms, which allow for more flexible and expressive transformations of probability measures. \n\n\"Distance-based tree sliced Wasserstein distance\" [TCNLN '25]", "essential_references_not_discussed": "I do not know of any essential references which are not discussed.", "other_strengths_and_weaknesses": "Strengths: The authors present a novel extension of the previous tree-sliced Wasserstein distance to use non-linear projections. These new tree-sliced Wasserstein distance variants may help better encode topological information than tree-sliced Wasserstein with linear projection. They also present extensive experiments that highlight the strength of their new tree-sliced Wasserstein variants. \n\nWeaknesses: Given that there are so many variants of TSW, the authors could maybe motivate a bit more why/when CircularTSW or SpatialTSW will outperform SW or TSW with linear projection.", "other_comments_or_suggestions": "You use the abbreviation for Circular Radon Transform on Systems of Lines (CRTSL) but you forgot to add the parenthetical defining the abbreviation above Eq. 10.", "questions_for_authors": "1. Is there any intuition for when each tree-sliced Wasserstein distance should be used? It seems now that is a large zoo of tree-sliced Wasserstein distances but I am unsure when I should use one tree-sliced Wasserstein distance over another.\n\n2. How should I select $r$ for CircularTSW? How does the choice $r$ affect performance empirically? \n\n3. It seems like this tree-sliced Wasserstein distance on systems of lines can be very general. Is there a way to define a general tree-sliced framework that will include both CRT and SRT?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zuUi8f1Jh3", "reviewer": "Reviewer_zCq9", "summary": "This paper extends the Tree-Sliced Wasserstein (TSW) distance, an alternative to the Sliced Wasserstein (SW) distance that leverages tree-based metric spaces, by allowing the use of nonlinear projections. More precisely, the authors explore generalized Radon transforms (previously used in existing SW variants, such as Generalized SW by Kolouri et al. (2019) or Augmented SW by Chen et al. (2022)), and analyze how these can be integrated into TSW while preserving injectivity and invariance properties. This analysis leads to the definitions of two instances of TSW, called \"Circular Tree-Sliced Wasserstein Distance\" and \"Spatial Tree-Sliced Wasserstein Distance\". Finally, they apply these new metrics to generative models, gradient flows and self-supervised learning, to compare the result quality and computational efficiency against SW and variants.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We direct the Reviewer to Table R1-4 and Figure R1 at https://sites.google.com/view/nonlinear-tsw-2.\n\n**Q1. Explicitly highlight ... to prior work.**\n\n**Answer Q1.** The key technical challenge in developing this approach lies in proving the injectivity of the proposed Radon Transforms, which ensures that the resulting distances qualify as proper metrics. This differs from prior work on TSW, such as [2]. Compared to SW approaches—including those that also employ nonlinear projection techniques like [1]—the difference is even more pronounced due to the introduction of splitting maps $\\alpha$.\n\nGiven these factors, we believe our approach to addressing these challenges is non-trivial.\n\n**Q2. The notions of well-definedness ... Radon Transform variants\".**\n\n**Answer Q2.** The two mentioned properties above hold under certain assumptions on the functions $g$ and $h$, such as continuity, smoothness, and injectivity. These conditions have been discussed in previous works (see lines 134–144).\n\nIn our work, we do not elaborate on these general assumptions; instead, we provide specific types of functions that satisfy the necessary conditions for our method. A detailed discussion of why these functions meet the required properties are presented Section 4.2 and Appendix B.\n\nInjectivity is typically required for variants of the Radon Transform because it ensures that the derived metrics are proper metrics rather than pseudo-metrics.\n\n**Q3. In what sense ... Radon Transform?**\n\n**Q4. Why does $\\alpha$ ... being referred to**\n\n**Answer Q3+Q4.** Previous works on TSW [2, 3, 4] show that replacing lines in the SW framework with tree structures via the splitting mechanism $\\alpha$ consistently enhances performance, even with the same number of projections.\n\nNotably, $\\alpha$ is independent of SW components, suggesting that other SW improvements could be integrated into the tree-sliced framework. However, verifying well-definedness and injectivity in this new setting requires novel analytical approaches.\n\n**Q5. Design a simpler ... clearly demonstrated.**\n\n**Answer Q5.**   The non-linearity in SpatialTSW complicates the design of settings where it outperforms other variants. Empirically, SpatialTSW matches Db-TSW (see Tables 1, 2, and R1), suggesting it as a drop-in replacement for Db-TSW with potential performance gains.\n\nWe provide intuition for when CircularTSW and CircularTSW$_{r=0}$ may outperform other variants. Since these distances rely on the $L_2$ norm for projection, they are likely to excel when the $L_2$ norms of the data are diversely distributed. We validate this advantage over Db-TSW and SpatialTSW in Table R2 on a synthetic dataset. The improvement is more pronounced for high-dimensional data (large $d$), indicating that variants using circular defining functions perform well while others struggle in such settings.\n\n**Q6. Across all experiments, ... image generation quality.**\n\n**Answer Q6.** In our paper, uncertainty quantification for Table 2 is provided in the Appendix (see Table 7). We also include uncertainty quantification for the Diffusion Model and SWAE experiments in Tables R3 and R4, respectively. \n\nWe provide a new qualitative result for Point Cloud Gradient Flow, visualizing the faster convergence of SpatialTSW in Figure R1.\n\n**Q7. The authors focus on the first-order ... and well-definedness.**\n\n**Answer Q7.** For $p>1$, the proposed approach can be extended. However, the Tree-Wasserstein distance with $p>1$ lacks a closed-form solution (see [5]). A meaningful alternative is provided by Sobolev Transport (ST) [6], which offers a closed-form solution and has been applied in the tree-sliced framework, as discussed in [3, Eq. (15)].\n\nAlthough TSW works such as [2] and [4] do not explicitly address this aspect, their implementations support arbitrary $p>1$. We omit the complex ST literature to reduce presentation complexity while keeping implementation flexible, and thus chose not to include it.\n\nDue to space constraints, we strongly encourage the Reviewer to refer to the ST literature [6]. Based on its properties, our extension to the $p>1$ case satisfies the theoretical guarantees previously discussed.\n\n**Q8. The results in Tables 1 and 2 ... is negligible**\n\n**Answer Q8.** We acknowledge that the reduction in computation time in Table 2 is marginal, as this experiment uses a toy synthetic dataset. However, in Table 1, when applied to real-world Diffusion Models, CircularTSW$_{r=0}$ reduces the total training time to 105.5 hours compared to 131 hours for Db-TSW over 1800 epochs, saving 25.5 hours of computation time.\n\n---\nWe thank the Reviewer for the constructive feedback, as well as for pointing out typos and missing references, which we will address. If our clarifications are satisfactory, we kindly ask you to consider raising the score. We are happy to address any further concerns in the next discussion stage.\n\n--- \n*References.* Kindly refer to **References** in our response to Reviewer jNth.", "claims_and_evidence": "The theoretical and methodological contributions of this paper seem sound to me, as they naturally extend related work by combining results from SW based on nonlinear projections and tree-based SW. That said, some imprecise points remain and could be addressed (see section 'Theoretical Claims').\n\nMy biggest concern is with the empirical analysis and conclusions. In my opinion, the authors make overclaims that are not properly justified by solid and convincing evidence: instead of providing a more nuanced description of the obtained results, they draw conclusions that seem overstated to me. See my detailed comments in \"Methods and Evaluation Criteria\" and \"Experimental Designs or Analyses\".", "methods_and_evaluation_criteria": "The experiments in this paper involve incorporating different metrics into existing black-box machine learning pipelines, whose behavior is hard to interpret. As a result, directly comparing the performance of different variants is challenging, and the analysis relies solely on quantitative scores such as FID, which I find insufficient to illustrate the authors' strong conclusions.\n\nI think the paper would be significantly more convincing if the authors identified a simpler and more interpretable setting where their proposed non-linear tree SW metrics demonstrably capture relevant data features more efficiently. Such a well-chosen controlled experiment could provide more insights into the advantages of their approach, that nicely complement the current empirical analysis.", "theoretical_claims": "I skimmed through the proofs in the supplementary document. The theoretical claims appear to be adaptations of results from the literature on generalized/augmented SW and tree SW, with these extensions facilitated by the linear operations in the tree structure.\n\nThe authors focus on the first-order ($p = 1$) case, justifying this choice by stating: \"For simplicity, the focus is on measures with a finite first moment, while measures with a finite $p$th-moment are treated analogously.\" (l.85-88) However, the treatment of $p>1$ is not as trivial as suggested when establishing the metric properties, injectivity, and well-definedness.", "experimental_designs_or_analyses": "The authors claim that the proposed tree-sliced-Wasserstein distances \"consistently outperform state-of-the-art Sliced Wasserstein and Tree-Sliced Wasserstein methods across various tasks\" and support this point with a series of experiments (Section 6 and supplementary doc). However, I find the empirical results not convincing enough:\n\n- Across all experiments, the improvement in precision appears marginal (see Tables 1-5), and there is no uncertainty quantification. Additionally, the only qualitative results available (Figure 7 in the supplementary document) do not clearly demonstrate the advantages of the proposed tree-sliced metrics in terms of image generation quality.\n\n- The authors state that their approach \"maintains comparable or improved runtime efficiency\" and specifically, they report that \"CircularTSW-DD and CircularTSW reduce training time relative to Db-TSW-DD⊥ by 10% and 19%, respectively\" and \"CircularTSW and CircularTSWr=0 are approximately 5% and 16% faster than vanilla SW, respectively\". While these claims are accurate, the results in Tables 1 and 2 reveal the reduction in computation time is negligible: CircularTSW and CircularTSWr=0 take 0.0017s and 0.0015s per iteration, respectively, compared to 0.0018s for SW. Furthermore, the results lack consistency across experiments: in Table 1 and Table 5, vanilla SW is the first or second fastest method.", "relation_to_broader_scientific_literature": "The key contribution of this paper is the combination of non-linear projections in sliced optimal transport with the tree structure. Both the theoretical results and experimental designs build on existing literature in this area, and the technical challenges involved in integrating these prior works are not sufficiently emphasized.", "essential_references_not_discussed": "The related works is adequately discussed. Some references that are relevant and seem to be missing are: \n- \"Parallelly Sliced Optimal Transport on Spheres and on the Rotation Group\", M. Quellmalz, L. Buecher, G. Steidl (2024)\n- \"Sliced Optimal Transport on the Sphere\", M. Quellmalz, R. Beinert, G. Steidl (2023)", "other_strengths_and_weaknesses": "I find the paper lacks clarity and rigor: \n- Several key definitions are missing from the main text, making it difficult to read and understand (see \"Other comments or suggestions\" for details).\n- Some strong claims about important aspects of the method lack proper and precise justification. For instance, the notions of well-definedness and injectivity (l.134) are not properly defined, and it is unclear why \"injectivity is typically required for Radon Transform variants\". Remark 3.2 presents intriguing, non-trivial information, without supporting results or references: why does $\\alpha$ induce a tradeoff between effectiveness and theoretical guarantees, , and what kind of guarantees are being referred to? In what sense \"the distances derived from RTSL surpass those obtained from the original Radon Transform\"?", "other_comments_or_suggestions": "- l.98: the definition of the set $L^1(\\cdot)$ is missing\n- l.107: missing definition for $\\mathcal{U}(\\mathbb{S}^{d-1})$\n- Missing definition for $\\theta_\\sharp$ (Section 2)\n- l.102 : $\\mathbb{R}_{\\geq 0}$ can be written $\\mathbb{R}_+$\n- Section 3.2: all the notations given in Section A to define tree-sliced SW should be recalled in that section.\n- Equation (14): the operation $gy,g\\mathcal{L}$ is unclear in the main text: equation (66) should appear earlier.\n- l.709: \"Sysmtes\" \n- Section 6.2: \"Inspised\"", "questions_for_authors": "Given my comments above, the proposed nonlinear tree SW variants seem to be incremental extensions of existing methods. This lack of novelty, in some sense, would not be problematic if the contributions yield clear advantages on other aspects, for instance here in terms of practical performance. However, in my opinion, this is not sufficiently achieved in the current work.\n\nTherefore, I would appreciate if the authors could:\n- Explicitly highlight the technical challenges involved in developing this approach or in establishing its theoretical guarantees, particularly in comparison to prior work.\n- Design a simpler, more interpretable experiment (e.g.,  based on synthetic data) where the advantages of their proposed methods are more clearly demonstrated.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BhnbImWefY", "reviewer": "Reviewer_jNth", "summary": "This work proposes to extend the Tree Sliced-Wasserstein distances, defined using linear projections on system of lines, by using nonlinear projections instead. The authors study the use of two different non linear projections: circular projections and spatial projections. They also propose to use a spatial projection on the sphere. Then, they introduce the associated Radon transform on system of lines, show that they are injective, define the resulting Tree-Sliced Wasserstein distances and show that these are well distances. Finally, they benchmark these distances on several applications such as generative modeling with Denoising Diffusion GANs and gradient flows for the Euclidean non linear projections, and gradient flows, self-supervised learning and Sliced-Wasserstein autoencoders for the spherical version.\n\n## Update after rebuttal\n\nI maintain my positive score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We direct the Reviewer to Table R1-2 at https://sites.google.com/view/nonlinear-tsw.\n\n**Q1. The paper feels a bit incremental, but there are lots of results, which compensate.**\n\n**It is not really stated when one would prefer one type of non linear projection compared to another.**\n\n**Answer Q1.**  We kindly refer the Reviewer to our response to **Q1+W1** in Reviewer m8Ge's review.\n\n**Q2. You are proposing new choices for the polynomials $h$: how does it compare to the construction proposed in [1]?**\n\n**Answer Q2.** In [1], the author proposed using homogeneous polynomials of odd degree for the mapping $h: \\mathbb{R}^d \\to \\mathbb{R}^{d_\\theta}$, where the output dimension $d_\\theta$ grows exponentially with $d$. For large $d$, such as in our Diffusion Experiment where $d \\approx 3000$, this results in $d_\\theta \\approx 4.5 \\times 10^9$ (see line 263-270), making the approach computationally impossible. In contrast, our proposed mapping, defined as $h(x) = (f_1(x), \\ldots, f_d(x))$ with $f_i(x) = x_i + x_i^3$, maintains the same output dimension as the input ($d_\\theta = d$) and introduces a trivial computational cost while still ensuring a non-linear projection.\n\n**Q3. How the trees are constructed? I am not sure it is precised (but I may have missed it)**\n\n**Answer Q3.** The tree structure used in our paper is the concurrent-lines structure introduced in [2]. We will make this explicit in the revised version of the paper.\n\n**Q4. The background is not very clear, and lot of important things to understand the paper are in Appendix. Also, more figures could help understanding the constructions.**\n\n**Answer Q4.**\n\nWe acknowledge that the background for the tree-sliced approach can be dense, as it builds upon the foundations laid in [2] and [4]. We thank the Reviewer for pointing this out and will revise the paper to include additional figures to improve readability.\n\n**Q5. In Figure 2, the result are given for $n$ up to $n=500$, which is rather small for sliced settings.**\n\nIn Figure 2, we use $n$ up to $500$ because Diffusion Model experiments typically train with batch sizes $n < 500$. Additionally, we set $d = 3000$, $L = 2500$, and $k = 4$, aligning with the practical settings of the Diffusion Model experiment. We also provide runtime and memory analysis for our Tree-Sliced Wasserstein variants in Appendix D.3, with $n$ up to $50000$.\n\n**Q6. In Table 2 and 3, this is basically the same experiment, but the results are not provided in the same way ($W_2 \\text{ versus } \\log_2 W_2$).**\n\nTable 2 is about Gradient Flow on Euclidean space $\\mathbb{R}^d$, while Table 3 is about Gradient Flow on the sphere $\\mathbb{S}^d$. Therefore, the choice of experimental setting and baselines are different. Note that Table 2 follows the setting of [2], while Table 3 follows the setting of [4].\n\n---\nWe thank the Reviewer for the constructive feedback, as well as for pointing out typos and missing references. We will address them accordingly. If the Reviewer finds our clarifications satisfactory, we kindly ask you to consider raising the score. We would be happy to address any further concerns during the next stage of the discussion.\n\n---\n*References.*\n\n[1] Soheil Kolouri et al., Generalized Sliced Wasserstein Distances. NeurIPS 2019.\n\n[2] Hoang Tran et al., Distance-Based Tree-Sliced Wasserstein Distance, ICLR 2025.\n\n[3] Hoang Tran et al., Tree-Sliced Wasserstein Distance on a System of Lines.\n\n[4] Hoang Tran et al., Spherical Tree-Sliced Wasserstein Distance, ICLR 2025.\n\n[5] Tam Le et al., Tree-Sliced Variants of Wasserstein Distances, NeurIPS 2019.\n\n[6] Tam Le et al., Sobolev Transport: A Scalable Metric for Probability Measures with Graph Metrics, AISTATS 2022.\n\n[7] Makoto Yamada et al., Approximating 1-Wasserstein Distance with Trees, TMLR 2022.", "claims_and_evidence": "The claims made are well supported. The constructions proposed are well justified, and the proofs of the claims such as distance properties are provided.", "methods_and_evaluation_criteria": "The evaluation criteria to compare the new distances with previously proposed distances make sense.", "theoretical_claims": "The proofs seem to be correct.", "experimental_designs_or_analyses": "The experimental designs are good.", "relation_to_broader_scientific_literature": "The key contribution of this paper is to extend the Tree Sliced-Wasserstein distance proposed in [1, 2] using non linear projections. Non linear projections were first used in [1] for the Sliced-Wasserstein distance. In particular, they leverage results from [2] to show that the proposed constructions are well distances.\n\nThe second key contribution is to extend the spherical Tree Sliced-Wasserstein distance proposed in [4] with non linear projections.\n\n\n[1] Tran, V.-H., Pham, T., Tran, T., Le, T., and Nguyen, T. M. Tree-sliced wasserstein distance on a system of lines. arXiv preprint arXiv:2406.13725, 2024.\n\n[2] Tran, H. V., Nguyen-Nhat, M.-K., Pham, H. T., Chu, T., Le, T., and Nguyen, T. M. Distance-based tree-sliced wasserstein distance. In The Thirteenth International Conference on Learning Representations, 2025.\n\n[3] Kolouri, S., Nadjahi, K., Simsekli, U., Badeau, R., and Rohde, G. Generalized sliced wasserstein distances. Advances in neural information processing systems, 32, 2019.\n\n[4] Tran, H. V., Chu, T., Nguyen-Nhat, M.-K., Pham, H. T., Le, T., and Nguyen, T. M. Spherical tree-sliced wasserstein distance. In The Thirteenth International Conference on Learning Representations, 2025.", "essential_references_not_discussed": "All the essential references seem to be discussed.", "other_strengths_and_weaknesses": "This paper introduces on one hand non linear projections to extend the Euclidean and Spherical Tree Sliced-Wasserstein distances, which are natural things to study. This paper is doing it well as it provides justifications to all the choices. However, the background is hard to follow, as it relies a lot on the previous papers [1,2,3], with a brief description in Appendix. This is the main limit that I found with the current state of this work.\n\n**Strengths**:\n- Provide new (Spherical) Tree SW distances using non linear projections, which are well distances. In particular, the Spatial Spherical TSW distance use a new projection.\n- Several experiments showing benefits compared to other distances\n\n**Weaknesses**:\n- The paper feels a bit incremental, but there are lots of results, which compensate.\n- The background is not very clear, and lot of important things to understand the paper are in Appendix. Also, more figures could help understanding the constructions.\n- It is not really stated when one would prefer one type of non linear projection compared to another.\n\n[1] Tran, V.-H., Pham, T., Tran, T., Le, T., and Nguyen, T. M. Tree-sliced wasserstein distance on a system of lines. arXiv preprint arXiv:2406.13725, 2024.\n\n[2] Tran, H. V., Nguyen-Nhat, M.-K., Pham, H. T., Chu, T., Le, T., and Nguyen, T. M. Distance-based tree-sliced wasserstein distance. In The Thirteenth International Conference on Learning Representations, 2025.\n\n[3] Tran, H. V., Chu, T., Nguyen-Nhat, M.-K., Pham, H. T., Le, T., and Nguyen, T. M. Spherical tree-sliced wasserstein distance. In The Thirteenth International Conference on Learning Representations, 2025.", "other_comments_or_suggestions": "I would suggest to try to improve the introduction of the background on the Tree SW distances, and to add Figures to better understand how the projections work. For instance, in Section 3.2, the Radon transforms on System of Lines are introduced, but all notations do not seem to be introduced (e.g. it is not clear directly what is $\\mathbb{L}_k^d$).\n\nIn Figure 1, it is not directly clear what is projected, what is $x_i$... It would be better to add a legend on the Figure with labels on the points projected and the center $x_i$. Also, it would maybe help to see a comparison with $r>0$.\n\nIn Section 5, it is stated that explanations are provided for the choice of the projections, and why they lead to more efficient metrics, but it is not very clear where it is stated after.\n\nIn Figure 2, the result are given for $n$ up to $n=500$, which is rather small for sliced settings.\n\nIn Table 2 and 3, this is basically the same experiment, but the results are not provided in the same way ($W_2$ versus $\\log_2 W_2$).\n\nTypos:\n- Line 204, 1st column: \"$R^n$\"\n- Line 215, 2nd column: \"We also examine different choices of functions that define the nonlinear projections explain why certain choices lead to more efficient metrics.\"\n- Line 293, 1st column: \"$f_i(x)=x_i+x_i^3$\"\n- Line 383, 2nd column: \"Inspised\"", "questions_for_authors": "1. You are proposing new choices for the polynomials $h$: how does it compare to the construction proposed in [1]?\n2. How the trees are constructed? I am not sure it is precised (but I may have missed it)\n\n[1] Kolouri, S., Nadjahi, K., Simsekli, U., Badeau, R., and Rohde, G. Generalized sliced wasserstein distances. Advances in neural information processing systems, 32, 2019.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xrdtxdo6g8", "reviewer": "Reviewer_kfHL", "summary": "The authors introduce the following: \n1. Generalized Radon transforms in the system of lines.   \n- These transforms extend the concept of the Radon transform by incorporating systems of lines and allowing for nonlinear projections, which improve the flexibility and applicability of the SW distance.\n2. Generalized Tree-sliced Wasserstein distance. \n- The authors develop two specific variants of TSW—the Circular Tree-Sliced Wasserstein distance and the Spatial Tree-Sliced Wasserstein distance, which offer more efficient metrics for measures on both Euclidean spaces and spheres.\n\n3. Applied generalized tree-sliced Wasserstein distance in generative model and gradient flow problems.\nIn addition, the new distance is applied for spherical data (Gradient flow and self-supervised learning).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We direct the Reviewer to Tables R1-2 and Figure R1, available at https://sites.google.com/view/nonlinear-tsw-4.\n\n**Q1. It seems that any continuous injective function $h$ can be used to define the generalized Radon transform. However, the criteria for selecting a suitable $h$ for different datasets or tasks remain unclear.**\n\n**Answer Q1.** We kindly refer the Reviewer to our response to **Q1+W1** in Reviewer m8Ge's review.\n\n**Q2. In Section 5.2, the authors explain that $h(x)$ can be determined by a neural network, but the structure and size of the neural network are not specified.**\n\n**When $h$ is implemented as a neural network, the cost of optimizing its parameters is not included in the computational complexity analysis in Section 5.3. This suggests that the actual complexity should be higher than the proposed value.**\n\n**Answer Q2.** A potential neural network for implementing $h(x)$ can be inspired by [6], where the neural network is a single feedforward layer that maps $\\mathbb{R}^d \\to \\mathbb{R}^d$. Such a layer would have $d \\times d$ weights and $d$ biases, introducing $O(d^2)$ parameters. If $h$ is implemented as a neural network, the cost of optimizing its parameters would scale the overall complexity linearly by the number of optimization steps. Incorporating learnable parameters into $h$ could enhance performance by better adapting to the data distribution, but it substantially increases the computational cost.\n\nIn our work, efficiency is a major concern, so we opted for a simple mapping to avoid this overhead. Specifically, for SpatialTSW, we use the mapping $h(x) = (f_1(x), \\ldots, f_d(x))$, where $f_i(x) = x_i + x_i^3$. This mapping costs $O(n \\cdot d_{\\theta})$, which is trivial compared to other steps like projection or sorting.\n\n**Q3. What is the reasoning behind the significant reduction in computational cost when using the Circular Radon Transformed Sliced-Wasserstein distance?**\n\n**Additionally, in Figure 2, the computational cost of circular TSW appears to increase significantly when $n \\geq 400$, but no explanation for this behavior is provided.**\n\n**Answer Q3.** The significant reduction in computational cost when using the Circular Sliced Wasserstein variants arises from the efficiency of computing $L_2$ norms compared to inner products. This advantage is illustrated in Figure R1.\n\nAs shown in Figure 2, the computational cost of $CircularTSW_{r=0}$ increases when $n \\geq 400$. This behavior is attributed to the cost of computing the splitting map $\\alpha$, which remains $\\mathcal{O}(L k n d_{\\theta})$ — the same complexity as other tree-sliced variants. While $CircularTSW_{r=0}$ is theoretically and empirically more efficient in the projection and sorting steps, the cost of computing $\\alpha$ becomes dominant as the number of support points $n$ grows.\n\nIn practice, $CircularTSW_{r=0}$ is slower than standard SW when the number of supports increases significantly since it involves the additional computation of the splitting map $\\alpha$, whereas SW does not. However, in a large-scale experiment such as a diffusion model, $CircularTSW_{r=0}$ increases training time by only $12$\\% compared to SW while substantially improving the FID from $3.64$ to $2.48$, highlighting a favorable practical trade-off. Additionally, $CircularTSW_{r=0}$ outperforms the best existing tree-sliced method, Db-TSW, in training time and FID.\n\n**Q4. The Circular Radon Transform (Equation (9)) is defined in $\\mathbb{R}^d$ rather than in a circular space. Could you clarify why it is referred to as the Circular Radon Transform?** \n\n**Answer Q4.**  The term circular refers to the circular defining function used in the Radon Transform, rather than the domain of the space itself. This naming convention is consistent with several prior works (see [4], [5]).\n\nIn contrast, for distances involving measures defined on the sphere, the term spherical is typically used (see [1], [2], [3]).\n\n---\nWe thank the Reviewer for the constructive feedback, as well as for pointing out typos and missing references, which we will address.  If the Reviewer finds our clarifications satisfactory, we kindly ask you to consider raising the score. We would be happy to address any further concerns during the next stage of the discussion.\n\n---\n*References.*\n\n[1] Hoang Tran et al., Spherical Tree-Sliced Wasserstein Distance, ICLR 2025.\n\n[2] Clément Bonet et al., Spherical Sliced-Wasserstein, ICLR 2023.\n\n[3] Huy Tran et al., Stereographic Spherical Sliced Wasserstein Distances, ICML 2024.\n\n[4] Soheil Kolouri et al., Generalized Sliced Wasserstein Distances, NeurIPS 2019.\n\n[5] Gaik Ambartsoumian et al., On the injectivity of the circular Radon transform, Inverse Problems 21 (2005).\n\n[6] Xiongjie Chen et al., Augmented Sliced Wasserstein Distances, ICLR 2022.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "New concepts: \n- Generalized  Radon transform in line system (Section 3.1, Eq(9)) and related injective property (Theorem 4.2-4.4) \n- non-linear tree sliced Wasserstein distance (Eq (17), (18)) and metric property (Theorem 5.3)\n- Additional RSpatial Spherical Radon Transform on Spherical Trees", "experimental_designs_or_analyses": "Yes. I've checked the experiment design.", "relation_to_broader_scientific_literature": "**References:**\n\n Kolouri, S., Park, J., & Rohde, N. (2019). Generalized Sliced Wasserstein Distances. arXiv preprint [arXiv:1906.06962](https://arxiv.org/abs/1906.06962).\n\nBonet, C., Berg, P., Courty, N., Septier, F., & Drumetz, L. (2022).\nSpherical Sliced-Wasserstein.\narXiv:2206.08780 | PDF\nDOI: 10.48550/arXiv.2206.08780\n\nTran, H., Bai, Y., Kothapalli, A., Shahbazi, A., & Liu, X. (2024).\nStereographic Spherical Sliced Wasserstein Distances.\narXiv:2402.02345 | PDF\nDOI: 10.48550/arXiv.2402.02345\n\nLeluc, R., Dieuleveut, A., Portier, F., & Segers, J. (2024).\nSliced-Wasserstein Estimation with Spherical Harmonics as Control Variates.\narXiv:2402.01493 | PDF\nDOI: 10.48550/arXiv.2402.01493\n\nRelation: Previous works have explored the combination of the generalized Radon transform/spherical Radon transform with the Wasserstein distance. This paper examines the application of the generalized Radon transform/spherical Radon transform to a system and analyzes the corresponding Wasserstein distance.", "essential_references_not_discussed": "Bonet, C., Berg, P., Courty, N., Septier, F., & Drumetz, L. (2022).\nSpherical Sliced-Wasserstein.\narXiv:2206.08780 | PDF\nDOI: 10.48550/arXiv.2206.08780\n\nLeluc, R., Dieuleveut, A., Portier, F., & Segers, J. (2024).\nSliced-Wasserstein Estimation with Spherical Harmonics as Control Variates.\narXiv:2402.01493 | PDF\nDOI: 10.48550/arXiv.2402.01493", "other_strengths_and_weaknesses": "### Weaknesses:  \n\n1. It seems that any continuous injective function \\( h \\) can be used to define the generalized Radon transform. However, the criteria for selecting a suitable \\( h \\) for different datasets or tasks remain unclear.  \n   In Section 5.2, the authors explain that \\( h(x) \\) can be determined by a neural network, but the structure and size of the neural network are not specified.  \n\n2. When \\( h \\) is implemented as a neural network, the cost of optimizing its parameters is not included in the computational complexity analysis in Section 5.3. This suggests that the actual complexity should be higher than the proposed value.  \n   Additionally, in Figure 2, the computational cost of circular TSW appears to increase significantly when \\( n \\geq 400 \\), but no explanation for this behavior is provided.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. The Circular Radon Transform (Equation (9)) is defined in \\(\\mathbb{R}^d\\) rather than in a circular space. Could you clarify why it is referred to as the Circular Radon Transform?  \n   Additionally, what is the reasoning behind the significant reduction in computational cost when using the Circular Radon Transformed Sliced-Wasserstein distance?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}]

Author rebuttal:

Meta-review: This paper proposes a nonlinear extension of Tree-sliced Wasserstein distance on a system of lines. More specifically, the generalized Radon transform on systems of lines and the spatial Radon transform on systems of lines are proposed. Through experiments, it demonstrates that the proposed method works well on Gradient flow, SSL, and encoder tasks.

Note that, for the tree Wasserstein distance, the original work was introduced by Indyk et al. 2003. However, it is completely missing and this is not fair to the authors of the original work. For the final version, I strongly request adding detailed discussion with the previous work including Indyk et al. 2003.
[1] Piotr Indyk, Nitin Thaper. Fast image retrieval via embeddings, 2003.

Overall, this is a solid paper, and I vote for acceptance if the above concern is clearly addressed in the final version.
Decision: Accept (poster)