Mixed-Curvature Tree-Sliced Wasserstein Distance
Authors: Duy-Tung Pham
Viet-Hoang Tran
Thieu Vo
Tan Minh Nguyen
Venue: ICLR 2026 Poster (2026)
Keywords: mixed curvature space
sliced optimal transport
Primary area: learning on graphs and other geometries & topologies
TL;DR: We propose MC-TSW, a valid and efficient metric for comparing distributions in mixed-curvature spaces, showing superior performance over product-space and constant-curvature baselines.

Abstract: Mixed-curvature spaces have emerged as a powerful alternative to their Euclidean counterpart, enabling data representations better aligned with the intrinsic structure of complex datasets. However, comparing probability distributions in such spaces remains underdeveloped: existing measures such as KL divergence and Wasserstein either rely on strong assumptions on distributions or incur high computational costs. The Sliced-Wasserstein (SW) framework provides an alternative approach for constructing distributional distances; however, its reliance on one-dimensional projections limits its ability to capture the geometry of the ambient space. To address this limitation, the Tree-Sliced Wasserstein (TSW) framework employs tree structures as a richer projected space. Motivated by the intuition that such a space is particularly suitable for representing the geometric properties of mixed-curvature manifolds, we introduce the Mixed-Curvature Tree-Sliced Wasserstein (MC-TSW), a novel discrepancy measure that is computationally efficient while faithfully capturing both the topological and geometric structures of mixed-curvature spaces. Specifically, we introduce an adaptation of tree systems and Radon transform to mixed-curvature spaces, which yields a closed-form solution for the optimal transport problem on the tree system. We further provide theoretical analysis on the properties of the Radon transform and the MC-TSW distance. Experimental results demonstrate that MC-TSW improves distributional comparisons over product-space-based distance and line-based baselines, and that mixed-curvature representations consistently outperform constant-curvature alternatives, highlighting their importance for modeling complex datasets.

Published as a conference paper at ICLR 2026
MIXED-CURVATURE TREE-SLICED WASSERSTEIN DISTANCE
Duy-Tung Pham" FPT Software Al Center Vietnam tungpdl 06 fpt.com
Viet-Hoang Tran"
Department of Mathematics
National University Singapore
hoang. tranviet @u. nus edu
Thieu N. Vo
Tan M. Nguyen
Department of Mathematics National University of Singapore thi chieuvoθnus, edd.. sg
Department of Mathematics National University of Singapore tanmn @ nus. edu..g sg
ABSTRACT
Mixed-curvature spaces have emerged a powerful alternative to their Eu- clidean counterpart, enabling data representations better aligned with the intrin- sic structure of complex datasets. However, comparing probability distributions in such spaces remains underexplored: existing measures such KL diver- gence and Wasserstein either impose strong distributional assumptions or suf- fer from high computational costs. The Sliced-Wassersee (SW) framework of- fers a tractable alternative for defining distributional distances, but its reliance one-dimensional projections limits its ability capture the geometric struc- ture of the ambient space and the input distributions. The Tree-Sliced Wasser- stein (TSW) framework provides a principled solution to this limitation by em- ploying tree structures as richer projected space. Motivated by the intuition that such a space is particularly suitable for representing the geometric proper- ties of mixed-curvature manifolds, we introduce the Mixed-Curvature Tree-Slice FreceSliced Wasserstein (MCTSW), a novel discrepancy measure that computationally ef- ficient while faithfully capturing both the topological and geometric structures of mixed-curvature spaces. Specifically, we introduce an adaptation of tree systems and Radon transform to lution for the optimal transport problem on the tree system. We further provide theoretical analysis on the properties of the Radon transform and the MCTSW distance. Experimental results demonstrate that MCTSW improves distributional comparisons over product-spacc--aasee distance and line-based baselines, and that mixed-curvature representations have better performance over constant-curvature alternatives, highlighting their importance for modeling complex datasets.
spaces, which yields closed- -foem so-
INTRODUCTION
Euclidean geometry has long served the default framework for representing data in machine learning. However, the manifold hypothesis suggests that real-world data typically concentrates on lower-dimensional structures rather than uniformly occupying Euclidean space (Loaiza-Ganne et al., 2024; Gray, 2025). This observation motivates the exploration of alternative geometries, particularly spherical and hyperbolic spaces, which capture structural properties that Fuclidean ge- ometry cannot (Cho al., 2023; Davidson et al., 2018; Dai et al., 2021; Yuan et al., 2025; Park & Kwon, 2019; Chen et al., 2022a). Spherical spaces naturally accommodate directional or periodic data, such as text embeddings (Meng et al., 2019; Banerjee et al., 2022: Xu et al., 2023) panoramic imagery (Yoon et al., 2022; Sun et al., 2025; Park et al., 2025). In contrast, hyperbolic spaces are especially effective for representing hierarchical structures (Lin et al., 2024b; Chami et al., 2020) and graph-structured data (Zheng et al., 2025; Liu et al., 2019).
Co-first authors. Please correspond hoang.tranvicc.. u.nus.edu and tanmn @ nus.edu.sg

---

Published as a conference paper at ICLR 2026
Mixed-Curvatee Space. Recent research has begun combining these geometries under the frame- work of Mrixed-Curvaturr Spaces (MCS), defined as Cartesian products of Euclidean, spherical, and hyperbolic components. MCS provide enhanced flexibility for modeling complex datasets (Skopek et al., 2020; Gu et al., 2019). These hybrid spaces have already demonstrated promise across a broad spectrum of models, ranging from decision trees and random forests (Chlenski et al., 2025), variational autoencoders (Skopek et al., 2020), and graph neural networks (Sun et al., 2022; Wang al., 2021), to applications in continual learning (Lin et al., 2024a), combinatorial optimization (Liu et al., 2025), and physical sciences (Woodward et al., 2024).
Optimal Transport. Despite recent advances, the fundamental challenge of comparing probability distributions across heterogeneous spaces remains unresolved. Standard divergences such as KL divergence are inadequate in this setting, as they are not true metrics and fail to handle distributions with disjoint supports. Optimal Transport (OT) (Villani, 2008; Peyré et al., 2019) provides geomet- rically natural metric for comparing probability distributions and has received significant attention in recent years. However, OT suffers from supercubic complexity in the number of supports, which severely limits its scalability (Peyré et al., 2019). To address this challenge, several scalable vari- ants have been developed, including entropic regularization (Cuturi, 2013), minibatch OT (Fatras et al., 2020), low low-rank methods (Forrow et al., 2019; Altschuler et al., 2019; Scetbon et al., 2021), the Sliced Wasserstein distance (Rabin et al., 2012; Bonneel et al., 2015), Tree-Sliced Wasserstein distance (Tran et al., 2025a; 2024c), and Sobolev transport (Le et al., 2022: 2024). Within the Fu- clidean setting, related methods have also been successfully applied in diverse domains, including machine learning (Kolouri et al., 2018b; Li et al., 2024; Carpintero Perez et al., 2024; Chapel et al., 2025; Lee al., 2019; Shahbazi et al., 2025; Doan et al., 2021: Wang et al., 2024: Rowland al., 2019; Melnyk et al., 2024), statistics (Hu & Lin, 2025; Yi & Liu, 2023; Tran et al., 2021; Bonet al., 2022b; Meunier et al., 2022; Chen & Müller, 2024), computer vision (Arjovsky et al., 2017; Genest et al., 2024; Fujiwara et al., 2024; Li et al., 2022; Yin & Chua, 2024; Bai et al., 2023), and natural language processing (Adhya & Sanyal, 2025; Otao & Yamada, 2023; Luong et al., 2025; Naderializadeh et al., 2021; Dao al., 2024: Kusner et al., 2015). More general formulations such as Partial Transport (PT) and Unbalanced Optimal Transport (UOT) have likewise found impact- ful applications in these areas (Séjourné et al., 2022; Kondratyev al., 2016; Eyring et al., 2024; Gazdieva et al., 2024; Chizat & Bach, 2018; Demetei et al., 2022; Baltrušaitis et al., 2018; Tran al., 2025d).
(Tree) Sliced Wasserstein. Among OT methods, the Sliced Wasserstein (SW) distance has gained prominence by reducing complexity: it projects high-dimensional probability measures onto one- dimensional subspaces, where OT admits closed-form solutions (Rabin et al., 2012; Bonneel et al., 2015). This projection-basse
strategy has inspired variety of extensions, including accelerated sampling (Nadjahi et al., 2021; Nguyen et al., 2024a; 2021), informed slicing (Deshpande et al., 2019; Nguyen et al., 2024c: Tran et al., 2024b; Nguyen & Ho, 2023), and generalized integration domains (Kolouri et al., 2019; Quellmalz et al., 2023; Chen et al., 2022b; Bonet et al., 2023b; 2025; 2023c). Yet, the inherent limitation to one-dimensional slices often prevents SW from capturing the full geometric and topological complexity of distributions in curved spaces. The Tree-Sliced Wasserstein (TSW) framework addresses this limitation by replacing linear projections with tree metric systems (Le et al., 2019; Le & Nguyen, 2021; Tran al., 2025a; Indyk & Thaper, 2003; Takezawa et al., 2021; Lin et al., 2025; Tran et al., 2026b; 2025b; 2026a; 2024d). Unlike lines, tree structures provide richer geometry while still permitting closed-form OT solutions, enabling improved modeling of spatial and topological features. Despite these advances, the potential of TSW in mixed-curvature latent spaces remains largely unexplored, even though such settings require combining heterogeneous geometric properties.
Contributions. In this work, we introduce principled framework for comparing distributions in mixed-curvature spaces. Our main contributions are as follows:
1. In Section 3, we introduce Srteed-Curature Trees, a metric space that admits a tree metric and enables closed-form Wasserstein distances across heterogeneous geometries. We further develop a Radon-type transform on Mixed-Curvature Trees to facilitate efficient distributional comparison.
2. In Section 4, propose the Mixed-Curvate Tree-Sliced Wasserstein (MCTSW) distance. This formulation allows mass to be transported through a tree spanning multiple subspaces, thereby jointly respecting diverse geometric structures.
2

---

Published as a conference paper at ICLR 2026
Figure Spaces constant curvature, building blocks of mixed-curvature spaces. (Left) Euclidean space EE,, (middle) Hyperbolic spaces H2 (top) and P2 (bottom) with corresponding geodesics, and (right) Spherical spaces S2 (top) and D2 (bottom) with corresponding geodesics.
3. In Section 5, we validate our framework through experiments gradient flow, variational au- toencoders, and graph self-supervised learning, demonstrating the improvements over baselines and the advantages when modeling with mixed-curvature space.
Supplementary material includes a table of notation, detailed discussion on the properties of our Radon transform and MCTSW distance, implementation detail MCTSW, and additional experi- mental details.
2 BACKGROUND
We begin by presenting the necessary preliminarics on mixed-curvature spaces and the family of Sliced Wasserstein distance.
2.1 CONSTANT-CURVATION AND
SPACES
Constant-Curvattre Spaces. A d-dimensional constant curvature space (CCS) Cj. is a smooth Riemannian manifold with constant sectional curvature K. Depending the sign of K, this gives rise to hyperbolic (K < 0), Euclidean (K = 0), or spherical (K > 0) geometries.
Embedded models for CCS in 12d+1 include the Lorentz model (hyperboloid) for hyperbolic space,
H - -{ ( Id41) E Rd++ 27 xx
x241
1
21
0
K
and the sphere for spherical space,
SK - (FI,.... Ta+1) [Rd+1 : +22
1 K
22
(2)
At each point u C the tangent space T.CC is Fuclidean, and the exponential and logarithmic maps exp! T,CC - C% and log C* Tg C provide diffeomorphisms that link the manifold with its tangent spaces.
Following (Skopek et al., 2020: Gao et al., 2023; Ganea et al., 2018), we employ stereographic pro- jection to present unified representation of these geometries. For (1) x and inverse are
Rd+1 the projection
K
2y
pk (11, x T)T) = 1+ K
P (v)
(3)
V
K
1
K
1
K
Geometrically, this maps either the sphere or the hyperboloid onto the hyperplane x via projection from (-1,0...., The induced models are the Poincaré ball pk
(K 0) and the
3

---

Published as a conference paper at ICLR 2026
projected hypersphere D (K 0) in R Both conformal to Euclidean space, with conformal factor 2 = 2/(1 + (1)))
To define operators that could work with both P. ID and R. we adopt the gyrovector space formalism (Ungar, 2008), presented in Appendix A. This framework provides a smooth extension of Euclidean operations, recovering standard structures the flat limit (Skopek et al., 2020).
Wrapped Normal Distribution. A fundamental Gaussian-like distribution on these manifolds is the Wrapped Normal Distribution (WND) (Nagano et al., 2019; Skopek et al., 2020). Sampling proceeds as
N/ Σ) E ToM, u == p- (v) Mo = exp (w).
(4)
Mixed-Curvature Spaces. A Mixed-Curvature Space (MCS) is the Cartesian product
M = K2
(5)
j=1
where each component Ck2 is constant curvature space with curvature Kj-
For point u (uI,...; um) E M, tangent spaces and maps are defined componentwise:
T.M
11 Tu,, Cky²
Exp<(1)
(expli (91); :.... exp Km (4m))
(6)
Logu (x)
(logni ( log
Km (Im)).
The distance decomposes as (I, y)² Em dk, (Ir, vi) and the canonical measure factorizes (M x .... x Am) = 17 1'22, (Aj)-
Operations such Möbius addition and scalar multiplication are similarly componentwise ex- tended. This construction has proven valuable in representation learning (Gu al., 2019), generative modeling (Skopek et al., 2020), and continual learning (Gao et al., 2023).
2.2 THE SLICED VARIANTS OF WASSERSTEIN DISTANCE
Sliced Wasserstein Distance. Sliced Wasserstein distance define on the Radon transform (Helga- son, 2011) R: L (Rd) L (R gd-1) such that
[ [(=)((x- (z, (2, dz,
RS(2, 0)
where 8 denotes the Dirac delta function and E gd-1, For μt v E P(Rd) with densities fas fur the Sliced Wasserstein distance (SW) (Rabin al., 2012) is given
SW v) = ( W.(Rfa(,,), (f)()))) (o),
(8)
-.
where is the uniform probability measure on gd-1
Monte Carlo estimations with sampled directions 0; ~ (0) yields
SWZ(H,V) L Σ (.(Rf((,0)) Rf((-)))
Beyond Euclidean slicing, manifold-adapted variants leverage projections compatible with the derlying geometry. Notable extensions include spherical slicing (Bonet et al., 2023a; Tran et al., 2024a), hyperbolic slicing (Bonet et al., 2022a), and slicing in Cartan- Hadamarr manifolds (Bonet
4

---

Published as a conference paper at ICLR 2026
al., 2025). Sampling strategies have also been refined for efficiency and variance reduction (Nguyen et al., 2024b).
Tree-Sliced Wasserstein. The Tree-Sliced Wasserstein (TSW) distance generalizes classical sliced Wasserstein methods by projecting measures onto tree metrics structures formed from connected line segments rather than single projections (Indyk & Thaper, 2003; Tran et al., 2025a). This involves partitioning measures via a splitting map and computing Wasserstein distances along the induced tree system, where closed-form solutions for W are available. Unlike prior works that fix the metric space as tree metric (Le et al., 2019; Lin et al., 2025; Indyk & Thaper, 2003; Takezawa al., 2021). our approach follows the setting (Tran al., 2025a: 2024c), wherein the tree system is sampled dynamically. This design specifically accommodates measures with dynamic support, enabling greater flexibility and applicability.
TREE SYSTEMS IN MIXED-CURVATURE SPACES
In this section, introduce the construction of tree systems within mixed-curvature spaces, which provides the geometric foundation for our generalized Sliced Wasserstein framework. We subse- quently propose Radon-type transform operator defined on such tree systems.
3.1 MIXED-CURE CUREDCCURVATURE TREE SYSTEMS
Definition 3.1 (Geodesic Ray in M). Let M be mixed-curvature space. For a. E M with direction y defined as
y * x, the geodesic ray r originating
2 = [Exp, MM (L. Log (v))
(10)
150
Remark 3.2. Any pair (t, r2) with > uniquely corresponds a point in M. Thus, this paramet- ric representation establishes a natural correspondence between points in the mixed-curvature space and elements the tree system.
We now construct tree structure generated by k geodesic rays emanating from common root x E M. Let L1, y/....,3k E M be points such that
dm(x, = for all i 1,...,k,
(11)
where is small constant. This normalization ensures that the directions of the rays well-defined, thereby determining the family of rays {r >
Consider the disjoint union| _ r equipped with the equivalence relation defined by
(ta, (ty, 1 = t; t,
(12)
We define the mixed-curvante curvatuue tree as the quotient space
T MI DE
(13)
which has root æ and edges given by the rays {r J canonical measure T V..... is induced by the standard Borel measure on [0, oo) through the parametric representation of geodesic rays.
We endow T T...... with tree metric.
Definition 3.3 (Tree Metric on T). For points = (f,, r: and b = (ty,r distance is defined as
r2, in Th the
dr(a, = 11 4 if j. ti tittt if j.
(14)
Theorem 3.4 (Metric Structure of Mixed-Curvature Trees). The space (T (T.....k, dr) a metric space endowed with a tree metric.
The proof follows directly from standard arguments for tree metrics, as in (Tran et al., 2024c), verifying non-negativity, symmetry, the triangle inequality, and separation of points.
5

---

Published as a conference paper at ICLR 2026
Remark 3.5 (Geometric Interpretation). The metric dy encodes the tree structure within mixed- curvature geometry. Specifically, when two points lic on the same geodesic ray = j), their distance coincides with the ambient metric along that ray. Conversely, when two points lie different rays (i # f), their distance is computed by traveling through the common root 2, which links rays that may correspond to regions of differing eurvature.
3.2 RADON TRANSFORM ON MIXED-CURVATI TREES
To formally define the Radon transform, we first need to specify how integration is defined on trees. For tree Tr..ni the integral of function given by
(L=)do(>)== =
-E ES. f(t,r%) dt.
(15)
The associated space ('(T) is then understood as the collection of all real- functions on T whose absolute value admits finite integral with respect to this measure.
To define the Radon transform onto a mixed-curvature tree, two components are required:
1. A projection functionn that assigns each point z E MM coordinate on the rays T.
2. A splitting map .M - Ak-1, where Ak-1 is the (k 1)-dimensional probability simplex, specifying how the mass each point is distributed across the rays.
For the projection function, we assign each point the coordinate that equal to dz, the geodesic distance to root, for every ray. We define the splitting map via the softmax of geodesic distances:
alz, T) exp(-dM(2,r?)
(16)
where (z, r min dm (h, 2), and t denotes the complete geodesic line through x and yi- ker
Other design choices for projection function and splitting map discussed in Appendix F. Trees). Let TM denote the collection of all mixed-curvature trees with l rays in M. For a E C(M x T Ax-1), the Radon transform is the operator
Definition 3.6 (Radon Transform on
R" (MM) - II L'(T)
(17)
TETM
defined by
(R" (t. / f(z) (() -1: do(z),
(18)
for each T T.M, where d is the Dirac delta distribution.
A detailed discussion of the properties of this operator, including well-definedness, linearity, and injectivity, is provided in Appendix B.
MIXED CURVATURE TREE SLICED WASSERSTEIN DISTANCE
We ready to present novel distributional distance Mixed-Curvature space in this section. Definition 4.1 (Mixed Curvature Tree Sliced Wasserstein Distance). The Mixed Curvature Tree Sliced Wasserstein distance between two probability measures v, P(M) is defined
-frr War (R+v, R71) do(T).
MCTSW(P,
(19)
where a is a probability measure on the space T. mixed curvature trees with k edges, and W W denotes the I-Wasserstein distance the tree metric space (T, dd)).

---

Published as a conference paper at ICLR 2026
x
7 1.1
221
Figure 2: Illustration of the tree system (in bold) and the Radon transform onto the mixed-curvature tree E2 x D2 x PP2. The tree system consists of 5 rays, where each ray is represented by ray in one component and point in the others. The Radon transform projects the points and U2 onto the 7th geodesic at U1 and U2,i respectively.
The metricity of MCTSW followed by the definition and the injectivity of our Radon transform which we defered the discussion to Appendix C.
In practice the integral in Definition 4.1 is approximated by Monte Carlo estimator:
L Way, (R (.F,, R4,1),
MCTSWL(U,M)) = Σ
(20)
where (T.)_ independently sampled from the distribution
T.M
Closed-form formula of War.. Let T rooted tree (as graph) with non- non-negative edge lengths, and the ground metric dr, i.e., the length of the unique path between two nodes. Given two prob-
ability distributions and supported on nodes of T, the Wasserstein distance with ground metric dr has closed form as follows:
W Wdr,1(H,V) =E We (a([(p)(T()))))))))))))))))))))) ((T(v)))
-'we
(21)
where ve is the endpoint of edge e that farther away from the tree root, I'(ve) is the subtree of T rooted at Ue. and is the length of e. The derivation detail is provided in Appendix D.2.
Simplifications and Tree Sampling. enable parallel computation and simplify the model, we impose the following constraints:
1. Each edge differs from the root x exactly one component:
Bi : (y x)² + 0, and (ys - x>² = 0, vi Vi
2. All components have equal dimension d: d.
Under these assumptions, the space TM mixed curvature trees is in bijection with
M x [m]] (M.)*,
where Me Rd denotes the d-dimensional sphere of radius This identification endows T&M with a natural measure inherited from the product measure. Sampling proceeds by:
Sampling the root x E M according to a Wrapped Normal distribution.
2. Sampling the components for the edges from [m]]*
3. Sampling directions (ex + y1)..... (Ox t y&) E M.. Alternatively, directions could be sampled on the sphere Sd since each y is only required for computing the distance to the ray, and only the direction in the ambient space is needed.

---

Published as a conference paper at ICLR 2026
Time Complexity. The time complexity of the algorithm is OLLn log Ldmnk), where L is the number of trees used for Monte Carlo estimation, denotes the number of rays per tree, n is the number of support points, d denotes the dimension of each component, and m is the number of components.
The first term, (In log n),, arises from the sorting step, which is performed once per tree. Be- cause all coordinates along each ray are identical after projection onto the root, a single sorting operation suffices for the entire tree. The second term, O(Ldmnk), accounts for the bulk of the mixed-curvature operators. Moreover, MCTSW is naturally well-suited to parallelization, since most computations are independent across components, which allows efficient implementation on GPUs and distributed systems. The overall algorithm for MCTSW is detailed in Algorithm 1.
Algorithm 1: Mixed Curvature Tree Sliced Wasserstein distance
Input: M = x _ C. E P(M) as v(x) ΣØ= By 6(x a,) and
[(z) プド= Uy 6(x az);; number of trees L; number of rays per tree l;; splitting
maps at.
for 8= to do
Sample 21(1) (O
(1) Unif((S (
index
(f)
Unif[L,m]
x² + (o!
i = index; (E)
Di
x
otherwise
Vi
(0
R
O
Construct the mixed curvature tree Te = -T
21 (...))))))"
(2)
Compute Wdr, 1 (R4, v, R2 u) by Eq ??:
1 Σ t=1
Compute (P,SW(,,
Wdr, (R% R9;u);
I
return MCTSW(P,R) (v,))
EXPERIMENTS
We present experiments to demonstrate: i. the effectiveness of our proposed Sliced Wasserstein (MCTSW) distance, and ii. the advantages of modeling data in mixed-curvature spaces compared to constant-curvatur alternatives. We evaluate across three settings: gradient flow, variational autoencoders (VABs), and graph self-supervised learning. Experimental details and training time for these experiments are provided in Appendix E.
Tree-
5.1 GRADIENT FLOW
Our first experiment considers learning a target distribution v from a source distribution μ using Riemannian gradient descent on MCTSW (v, μ1). The target distribution v is constructed from 2,400 samples of a mixture of six wrapped-norma distributions (WNDs) arranged equidistantly from origin, while the source is a single WND centered at the origin.
We compare MCTSW against three baselines: SW SWamiient sliced Wasserstein (Bonneel et al., 2015) computed in the ambient space; ii Prod-TSW: product metric applying Tree-Sliced Wasser- stein independently on each factor space; iii. Prod-SW: product metric applying geodesic sliced Wasserstein in each factor space (Rabin al., 2012; Bonet et al., 2023b;a).
All methods are trained with butl-batch updates for 10,000 iterations, using 1,800 randomm projections per run. Table 1 shows that MCTSW achieves substantially lower geodesic log W₂ than other baselines. This suggests improved convergence behavior and more faithful recovery of the target distribution in mixed-curvature spaces.
8

---

Published as a conference paper at ICLR 2026
5.2 GRAPH SELF-SUPERVISED LEARNING
We further benchmark MCTSW in a graph self-supervised learning (SSL) setting, following (Liu al., 2024). Specifically, we learn node representations using loss function composed of an alignment term and a uniformity term:
L Lali + XLunis
where the uniformity loss is conventionally expressed as the W-aassesteein distance between hidden representations and Gaussian prior. In our approach, replace this term with MCTSW when modeling the latent space as an MCS manifold, obtaining MCTSW-SSGE, and compare against constant-curvature variants (E-, H-, and S- S-TSW-SSGE), which model the latent space in Euclidean, Hyperbolic, and Spherical spaces, respectively, while using MCTSW restricted to a single compo- nent as uniformity loss.
Table 2: Test Accuracy (1) on Cora.
Table 1: Learning a target mixture of 6 WNDs.
Method
Accuracy (1)
Method
log W2 (1)
SSGE
79.55 + 0.35
ammmieet
0.33
E-TSW-SSGE
77.85 + 0.32
Prod-TSW
0.34
H-TSW-SSGB
75.10 0.22
Prod-SW
0.31
S-TSW-SSGE
78.33 + 0.15
MCTSW (ours)
-3.65
MCTSW-SSGE 79.86 + 0.45
Obtaining the representation, we evaluate its quality through downstream classification accuracy using a random forest classifier on MCS (Chlenski et al., 2025). The results, shown in Table 2, demonstrate slight improvements over the baseline SSGE (Liu et al., 2024) on the Cora dataset (Kipf & Welling, 2017) when using MCTSW as the uniformity loss. Furthermore, modeling with an MCS latent space yields higher accuracy compared constant-curvature variants.
5.3 MIXED-CURVATURE VARIATIONAL AUTOENCODERS
Variational Autoencoders (VAEs) (Kingma & Welling, 2013) are probabilistic generative mod- els that learn continuous latent representations, enabling both faithful reconstructions and realis- tic data generation via sampling. Standard VAEs optimize a reconstruction loss combined with Kullback-Leibler (KL) divergence regularization, but subsequent work has investigated alternative divergences such as Wasserstein distances (Tolstikhin et al., 2018; Kolouri et al., 201 8a), well as non-Euclidean latent geometries including spherical (Davidson et al., 2018), hyperbolic (Bonet al., 2022a), and mixed-curvature spaces (MCS) (Skopek et al., 2020).
We evaluate our proposed MCTSW regularizer within a mixed-curvature latent space on CIFAR- 10 (Krizhevsky & Hinton, 2009), optimizing binary cross-entropy (BCE) reconstruction objec- tive. We compare against range of VAE baselines: the vanilla Euclidean VAE with KL diver- gence (Kingma& & Welling, 2013); SWAE (Kolouri al., 2018a) (Euclidean latent space with Sliced Wasserstein); S-V (Davidson et al., 2018) (spherical latent space with KL divergence); STSW. VAE (Tran et al., 2025c) (spherical latent space with tree-sliced Wasserstein); H- VAE (hyperbolic latent space with KI. divergence); HSW-VAE (Bonet et al., 2022a) (hyperbolic latent space with hyperbolic SW); and M-VAE (Skopek al., 2020) (MCS latent space with KI. divergence).
As shown in Table 3, incorporating MCTSW within MCS improves performance compared to i. the KL-based mixed-curvature baseline M-VAE, demonstrating the superiority of MCTSW over KL divergence for regularization, and ii constant-curvature KL- and SW-based variants, highlighting the advantages of modeling with MCS when combined with the MCTSW distance.
CONCLUSION
This paper introduced the Mixed-Cuue Mixed-Curvature Curvature Tree-Sliced Wasserstein Distance (MCTSW), novel measure leveraging tree systems to accurately represent both geometric and topological structures

---

Published as a conference paper at ICLR 2026
Table 3: Test BCE (1) on CIFARIO.
Latent space
Method
Regularizer
Test BCE (1)
VAE SWAE
-iv-eieeeeence
0.6423 + 0.0008
Euclidean
SW
0.6043 + 0.0005
S-VAE
KL-divergence
0,6285 = 0.0004
Spherical
STSW-VAE
STSW
0.6026 + 0.0009
H-VAE
KL-divergence
0,6402 + 0.0005
Hyperbolic
HSW VAE
HSW
0.6012 = 0.0006
M-VAE
KL-divergence
0.6419 + 0.0008
MCS
MCTSW-VAE
MCTSW
0.6000 + 0.0002
in mixed-curvature spaces. By adapting Radon transforms to tree- rre-based representations, MCTSW offers tractable and closed-form method for comparing probability distributions. Our theoretical analysis demonstrates key properties of the Radon transform and the MCTSW distance, ensuring the reliability for distributional comparisons. Empirical results in three main tasks: gradient flow, varia- tional auto-encoder, and graph self-supervised learning, show the advantages of MCTSW over base- lines methods and also illustrate the benefit of mixed-curvature latent representations over constant- curvature alternatives for modeling complex datasets. The method's 's design allows for efficient paral- lel computation, enhancing scalability; however, the computational overhead from gyrovector space operations poses challenges for runtime and numerical stability, which warrant further investigation.
ACKNOWLEDGMENTS
This research project is supported by the National Research Foundation Singapore under the AI Singapore Programme (AISG Award No: AISG2-TC-2023-012---5
This research project is supported by the Ministry of Education, Singapore, under the Academic Research Fund Tier (FY2023) (A-8002040-00-00)) A-8002039-00-00) This research / project is also supported by the NUS Presidential Young Professorship Award (A-0009807-01-00) and the NUS Artificial Intelli- gence Institute-Seed Funding -A-8003062-00-00,
Ethics Statement. Given the nature of the work, we do not foresee any negative societal and ethical impacts of our work.
Reproducibility Statement. Source codes for experiments are provided in the supplementary materials of the paper. The details of our experimental settings and computational infrastructure are given in Section 5 and Appendix E. All datasets that we used in the paper are published, and they are easy to access in the Internet.
LLM Usage Declaration. We use large language models (LLMs) for grammar checking and cor- rection.
REFERENCES
Suman Adhya and Debarshi Kumar Sanyal. S2WTM: Spherical sliced- Wasserstein autoencoder for topic modeling. In Proceedings of the 63rd Annual Meeting of the Association for Compu- tational Linguistics (Volume I: Long Papers), pp. 23211 23211-23225 Association for Computational Linguisties, July 2025. doi: 10.18653/v1/220 222/2/22222cooonnppo 1131.
Mark Agranovsky, Peter Kuchment, and Eric Todd Quinto. Range descriptions for the spherical mean radon transform. Journal of Functional Analysis, 248(2)3444-88) 2007. ISSN 0022-1236
Jason Altschuler. Francis Bach, Alessandro Rudi, and Jonathan Niles-Weed. Massively scalable sinkhorn distances via the nyström method. In Advances in Neural Information Processing Sys- tems, volume 32, 2019.
Yuri A. Antipov, Ricardo Estrada, and Boris Rubin. Inversion formulas for the spherical means in constant curvature spaces, 2011. URL https: // ///////atii/ rxiv. wtgg/bs///////////// aatpaa/./..11111 .992.
10

---

Published as a conference paper at ICLR 2026
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In Proceedings of the 34th International Conference on Machine Learning. volume 70 of Pro- ceedings of Machine Learning Research, pp. 223. PMLR, 2017.
Yikun Bai, Bernhard Schmitzer, Matthew Thorpe, and Soheil Kolouri. Sliced optimal partial trans- port. In 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp- 13681-13690, 2023. doi: 10. 00.1099CCVPR2272222000013133
Tadas Baltrušaitis, Chaitanya Ahuja, and Louis-Philippe Morency. Multimodal machine learning: A survey and taxonomy. IEEE transactions pattern analysis and machine intelligence, 41(2): 423 443, 2018.
Souvik Banerjee, Bamdev Mishra, Pratik Jawanpuria, and Manish Shrivastava Shrivastava. Gener- alised spherical text embedding. In Md. Shad Akhtar and Tanmoy Chakraborty (eds.), Proceed- ings of the 19th International Conference on Natural Language Processing (ICON), pp. 80-85, December 2022.
Clément Bonet, Laetitia Chapel, Lueas Drumetz, and Nicolas Courty. Hyperbolic wasserseens via geodesic and horospherical projections. Technical report, arXiv preprint arXiv:2211.1000 2022a. URL https: http///awwv..rraaaaaaeaaaaaaea/a ICML Workshop.
10066. Accepted at the TAG-ML 2023
Clément Bonet, Nicolas Courty, François Septier, and Lucas Drumetz. Efficient gradient flows in sliced-wasserstein space. Transactions on Machine Learning Research, 2022b. ISSN 2835-8856.
URL www.thtpp//// htnttts:://ppereveteet
ftw.e/wwwwee/fommmmmmmm id=Aui LNKmRvh.
Clément Bonet, Paul Berg, Nicolas Courty, François Septier, Lucas Drumetz, and Minh-Tan Pham. Spherical sliced-wasserstein In International Conference on Learning Representa- abs/
tions (ICLR), 2023a. doi: 10. 444455//AXT..226.... 087880 URL wtttts://aaxii..ttt 2206. 08780. Published as conference paper at ICLR 2023.
Clément Bonet, Laetitia Chapel, Lucas Drumetz, and Nicolas Courty. Hyperbolic sliced-Wasee Wasserstein via geodesic and horospherical projections. In Topological, Algebraic and Geometric Learning Workshops 2023, pp. 334-370. PMLR, 2023b.
Clément Bonet, Benoît Malézicux, Alain Rakotomamoniy, Lucas Drumetz, Thomas Moreau, Matthieu Kowalski, and Nicolas Courty. Sliced- Geeddd-Wssessttin on symmetric positive definite matri- ces for M/EEG signals. Proceedings of the 40th International Conference on Machine Learn- ing, volume 202 of Proceedings of Machine Learning Research, pp. 2777-2805. PMLR, 23-29 Jul 2023c.
Clément Bonet, Lucas Drumetz, and Nicolas Courty. wassert-nnn distances and flows on cartan-badamard manifolds. Journal of Machine Learning Research, 26(32):1-76, 2025. URL 0.4...039.pcc
https:// fmwr. wwtp//ppperrsse.eeeoeeeoe 6 wwweeerss/oooeeeeeeee
Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and Radon Wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision, 51:22 45, 2015.
Nicolas Boumal. An introduction to optimization on smooth manifolds. Cambridge University Press, 2023. doi: 10.1017/978 1009166164 URL https: : /wwww www.ooicoouuaaaa.........ccaaa.a net/book.
Raphaël Carpintero Perez, Sébastien Da Veiga, Josselin Garnier, and Brian Staber. Gaussian process regression with sliced Wasserstein Seisfeiler-Lehhmar graph kernels. In Proceedings of The 27th International Conference on Artificial Intelligence and Statisties, volume 238 of Proceedings of Machine Learning Research, pp. 1297-1305. PMLR, 02-04 May 2024.
Ines Chami, Albert Gu, Vaggos Chatziafratis, and Christopher Ré. From trees to continuous embed- dings and back: Hyperbolic hierarchical clustering. H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 15065-15076. 2020.
Lactitia Chapel, Romain Tavenard, and Samuel Vaiter. Differentiable generalized sliced wasserstein plans, 2025. URL https: /aw.xii..... org/ abw//bbs/555 22049.

---

Published as a conference paper at ICLR 2026
Han Chen and Hans-GGory Müller. Sliced wasserstein regression, 2024. URL https: https::/arxiv. axq/abs/2a 2306 6. 10601.
Tianyi Chen, Yunfei Zhang, Xiaoyang Huo, Si Wu, Yong Xu, and Hau San Wong. Spheriegan: Semi- supervised hyper-spherical generative adversarial networks for fine-grained image synthesis. In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 9991- 10000, 2022a.
Xiongjie Chen, Yongxin Yang, and Yunpeng Li. Augmented sliced wasserstein distances. In ternational Conference Learning Representations, 2022b. URL bittps://pperrvvtte net / forummidd iMqTLyEWnen Ewnoo.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing sys- tems, 31, 2018.
Philippe Chlenski, Quentin Chu, Raiyan R. Khan, Kaizhu Du, Antonio Khalil Moretti, and Itsik Pe'er. Mixed-curvature decision trees and random forests. In Forty-second International Confer- ence on Machine Learning, 2025.
Seunghyuk Cho, Juyong Lee, and Dongwoo Kim. Hyperbolic vae via latent gaussian distributions. In Advances in Neural Information Processing Systems, volume 36, pp. 569-588. Curran Asso- ciates, Inc., 2023.
CP Algorithm. Temary search. ternary_searee haarrchhhmm....... html.
https://pp-algorttame com www.ccwmmmmmctttcccccttc...teaet http:/uuummmmthheee
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013.
Shuyang Dai, Zhe Gan, Yu Cheng, Chenyang Tao, Lawrence Carin, and Jingjing Liu. APo-VAE. Text generation in hyperbolic space. In Proceedings of the 2021 Conference of the North Amer- ican Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 416-431, Association for Computational Linguistics, June 2021.
Viet Dao, Van-Cuong Pham, Quyen Tran, Thanh-Thiee Le, Linh Van Ngo. and Thien Huu Nguyen. Lifelong event detection via optimal transport. In Proceedings of the 2024 Conference on Empir- ical Methods in Natural Language Processing, pp. 12610-12621, Miami, Florida, USA, Novem- ber 2024. Association for Computational Linguistics.
Tim R. Davidson, Luca Falorsi, Nicola De Cao, Thomas Kipf, and Jakub M. Tomezak. Hyperspheri- cal variational auto-encoders. 34th Conference on Uncertainty in Artificial Intelligence (UAI-18), 2018.
Pinar Demetci, Rebecca Santorella. Manav Chakravarthy, Bjorn Sandstede, and Ritambhara Singh. Scotv2: Single-cell multiomic alignment with disproportionate cell-type representation. Journal 2022.
of Computational Biology,
Ishan Deshpande, Yuan- Ting Hu, Ruoyu Sun, Ayis Pyrros, Nasir Siddiqui, Sanmi Koyejo, Zhizhen Zhao, David Forsyth, and Alexander G. Schwing. Max-sliced wasserstein distance and its use for gans. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2019.
Khoa Doan, Yingjie Lao, and Ping Li. Backdoor attack with imperceptible input and latent modi- fication. In Advances in Neural Information Processing Systems, volume 34, pp. 19944-18957 Curran Associates, Inc., 2021.
Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian J Theis. Unbalancedness in neural Monge maps improves unpaired domain translation. In The Twelfth International Conference on Learning Representations, 2024. URL https: /poeereeview net. 1ttp//o
12

---

Published as a conference paper at ICLR 2026
Kilian Fatras, Younes Zine, Rémi Flamary, Remi Gribonval, and Nicolas Courty. Learning with minibatch wasserstein asymptotic and gradient properties. In Proceedings of the Twenty Third International Conference Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 2131-2141. PMLR, 26-28 Aug 2020.
Aden Forrow, Jan-Christian Hütter, Mor Nitzan, Philippe Rigollet. Geoffrey Schiebinger, and Jonathan Weed. Statistical optimal transport via factored couplings. Proceedings of the Twenty-Secome Second International Conference Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pp. 2454-2465. PMLR, 16-18 Apr 2019.
Haruo Fujiwara, Yusuke Mukuta, and Tatsuya Harada. Style-nerfnneff: 3d style transfer from style- aligned multi-view images. In SIGGRAPH Asia 2024 Conference Papers, 2024.
Octavian-Fugen Ganea, Gary Bécigneul, and Thomas Hofmann. Hyperbolic neural net- works. In Advances in Neural Information Processing Systems 31 (NeurIPS 2018), pp. 5350-5360, 2018. URL https: http://ptccceccccccceoe/o/c neurips httpppapp/2/1///eeeeee (Bab2addc899078000500
waeeeffaBllllllllo 122--14222-00 42- Abstract... html.
Zhi Gao, Chen Xu, Feng Li, Yunde Jia, Mehrtash Harandi, and Yuwei Wu. Exploring data geometry for continual learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 24325 24334, June 2023.
Milena Gazdieva, Arip Asadulaev, Evgeny Burnaev, and Alexander Korotin. Light unbalanced optimal transport. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https::///// http//ppppneeeeeeeeeeeeeeeeeeeee net forumm betrrmmid==coBae id=co8k2ws
Baptiste Genest, Nicolas Courty, and David Coeurjolly. Non-euclidean sliced optimal transport sampling, 2024. httLtttts//aaaett org/ absp/4bs/22222 16981.
Annie Gray. Statistical Exploration of the Manifold Hypothesis.
Phd the-
sis. University of Bristol, 2025. URL https://hdl
bf5e4bce- 2b0a-46b6----2 91 21166221ee55555155
Albert Gu, Frederic Sala, Beliz Günel, and Christopher Ré. Learning mixed-curvature represen-
tations products model spaces. International Conference on Learning Representations (ICLR), 2019. URL /ttpstts:://peeeeeeee http//ppneviee.eceommeeee .ntp//neeeeeeeeeeeee how.:///oooomm.dddeeeeiiee..
Sigurdur Helgason. The radon transform on R" In Integral Geometry and Radon Transforms, pp. 1- 62. Springer, New York, NY, 2011. ISBN 978 177-1-4419-6054-22 doi: 00.1007/978-1-4419--65555 1.
Xiaoyu Hu and Zhenhua Lin. Two-sample distribution tests in high dimensions via max-sliced wasserstein distance and bootstrapping. Biometrika, 112(2)aaaf000,, 01 2025. ISSN 1464-3510. doi: 10. 10.10933biomme/asaa000 URL https::/dss org / 1013/0io0oomeee asa f001.
Piotr Indyk and Nitin Thaper. Fast image retrieval via embeddings. Proceedings of the 3rd International Workshop Statistical and Computational Theories of Vision (SCTV 2003), 2003.
Svetlana Katok. Fuchsian Groups. The University of Chicago Press, 1992.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. 2013.
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference Learning Representations, 2017. URL https: / //openreview. net/ httt//orrmmii--sSuaaaaa
Soheil Kolouri, Phillip E. Pope, Charles E. Martin, and Gustavo K. Rohde. Sliced wasserseen autoencoder: An embarrassiggly simple generative model, 2018a. URL https: /axviv. org/abs//884 01 947.
13

---

Published as a conference paper at ICLR 2026
Soheil Kolouri, Gustavo K. Rohde, and Heiko Hoffmann. Sliced wasserstein distance for learning gaussian mixture models. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2018b.
Soheil Kolouri, Kimia Nadjahi, Umut Simsekli, Roland Badeau, and Gustavo Rohde. Generalized sliced wasserstein distances. In Advances in Neural Information Processing Systems, volume 32, 2019.
Stanislav Kondratyev, Léonard Monsaingeon, and Dmitry Vorotnikov. A fitness-driven cross- diffusion system from population dynamics as a gradient flow. Journal of Differential Equations, 26165):784-2808, 2016.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical Report 0, University of Toronto, Toronto, Ontario, 2009. URL https::/wwit /www.cs..www..... cs. eeronto.edu/-keee teaann/////tiii/eeeeeannggnn fettrrffttttteeeee -aaiuess-000--a 20.9-IR.pd.. pdf.
Matt Kusner, Yu Sun, Nicholas Kolkin, and Kilian Weinberger. From word embeddings doc- ument distances. In Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 957-966, Lille, France, 07-09 Jul 2015. PMLR.
Tam Le and Truyen Nguyen. Entropy partial transport with tree metrics: Theory and practice. In International Conference on Artificial Intelligence and Statistics, pp. 3835-3843. PMLR, 2021.
Tam Le, Makoto Yamada, Kenji Fukumizu, and Marco Cuturi. Tree-sliced variants of wasserstein distances. In Advances in Neural Information Processing Systems, volume 32, pp. 12283-12294, 2019.
Tam Le, Truyen Nguyen, Dinh Phung, and Viet Anh Nguyen. Sobolev transport: A scalable metric for probability measures with graph metrics. In Proceedings of The 25th International Confer- ence on Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pp. 9844 9868. PMLR, 28 30 Mar 2022.
Tam Le, Truyen Nguyen, and Kenji Fukumizu. Generalized sobolev transport for probability mea- sures on graph. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 26152-26174, PMLR, 21-27 Jul 2024.
Chen-Yo Yu Lee, Tanmay Batra, Mohammad Haris Baig, and Daniel Ulbricht. Sliced wasserstein discrepancy for unsupervised domain adaptation, 2019. URL hwtpt:://exxxi..oooo 1903. 04064.
Jie Li, Dan Xu, and Shaowen Yao. Sliced wasserstein distance for neural style transfer. Computers & Graphics, 002:89-98, 2022. ISSN 0097-8493. doi: https:///di..oo 10//1.0.100 (0016//.aa.2221.12222
Tao Li, Cheng Meng, Hongteng Xu, and Jun Yu. Hilbert curve projection distance for distribution comparison. IEEE Transactions on Puttern Analysis and Machine Intelligence, 46(7):4993-5007, 2024. doi: 10. 10.109/TPAM1..00.33367777
Ya-Wa Wei Eileen Lin, Ronald R. Coifmann Gal Mishne, and Ronen Talmon. Tres-wasserntein distance for high dimensional data with a latent feature hierarchy. In International Conference on Learning Representations (ICLR), 2025.
Zhanran Lin, Puheng Li, and Lei Wu. Exploring neural network landscapes: Star-shaped and geodesic connectivity. arXiv preprint arXiv:2404.06391 2024a.
Zhicheng Lin, HeGang Chen, Yuyin Lu, Yanghui Rao. Hao Xu, and Hanjiang Lai. Hierarchical topic modeling via contrastive learning and hyperbolic embedding. In Proceedings of the 2024 Joint International Conference on Computational Linguistics, Language Resources and Evalua- tion (LREC-COLING 2024), May 2024b.
Qi Liu, Maximilian Nickel, and Douwe Kiela. Hyperbolic graph neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché- Buc, Fox, and R. Garnett (eds.), Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
14

---

Published as a conference paper at ICLR 2026
Suyu Liu, Zhiguang Cao, Shanshan Feng, and Yew-Soon Ong. A paradigm for multi-task vehicle routing solver. In Forty- second International Conference on Ma- chine Learning, 2025.
based pre-training
Yunhui Liu, Tieke He, Tao Zheng, and Jianhua Zhao. Negative-free self-supervised gaussian em- bedding of graphs. Neural Networks, 2024.
Gabriel Baniza,Ganem,, Brendan Leigh Ross, Rasa Hosseinzadeh, Anthony L. Caterini, and Jesse C Cresswell. Deep generative models through the lens of the manifold hypothesis: A survey and new connections. Transactions on Machine Learning Research, 2024.
Manh Luong, Khai Nguyen, Dinh Phung, Gholamreza Haffrri, and Lizhen Qu. Unbiased sliced wasserstein kernels for high-quality audio captioning, 2025. URL https:: /ttp:///axvvvorr//// abs/2502. 05435.
Igor Melnyk, Youssef Mrouch, Brian Belgodere, Mattia Rigotti, Apoorva Nitsure, Mikhail Yurochkin, Kristjan Greenewald, Jiri Navratil, and Jarret Ross. Distributional preference align- ment of LLMSs via optimal transport. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https:/ /ttps:///pperevieww wet.eewwweeeeeeeeeee met/forum?id 2LotgfN6Ty.
Yu Meng, Jiaxin Huang, Guangyuan Wang, Chao Zhang, Honglei Zhuang, Lance Kaplan, and Jiawei Han. Spherical text embedding. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019, URL https::/procceediinn Files/papezs 210p///1100002222211111111111111 (at-8/99666ddc666---
neurips cc/paper_ pdf.
Dimitri Meunier, Massimiliano Pontil, and Carlo Ciliberto. Distribution regression with sliced Wasserstein kernels. Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference Ma- chine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 15501-15523. PMLR. 17-23 Jul 2022.
Navid Naderializadeh, Joseph F Comer, Reed Andrews, Heiko Hoffmann, and Soheil Kolouri. Pool- ing by slised-wasserstein embedding. In Advances in Neural Information Processing Systems, volume 34, pp. 3389-3400, 2021.
Kimia Nadjahi, Alain Durmus, Pierre E Jacob, Roland Badeau, and Umut Simsekli. Fast approxi- mation of the sissed-wasersteiin distance using concentration of random projections. In Advances in Neural Information Processing Systems, volume 34, pp. 12411-12424, 2021.
Yoshihiro Nagano, Shoichiro Yamaguchi, Yasuhiro Fujita, and Masanori Koyama. A wrapped nor- mal distribution hyperbolic space for gradient-based learning. In Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 4693-4702. PMLR, 2019.
Khai Nguyen and Nhat Ho. Energy-based based sliced wasserstein distance. Advances in Neural Information Processing Systems, volume 36, pp. 18046-18075, 2023.
Khai Nguyen, Nhat Ho, Tung Pham, and Hung Bui. Distributional sliced- wasserstein and applica- tions to generative modeling. In International Conference on Learning Representations, 2021. net/ /ooofooooo/ forum?1d-00
URL
Khai Nguyen, Nicola Bariletto, and Nhat Ho. Quasi-monte carlo for 3d sliced wasserstein. In The Tweifth International Conference on Learning Representations, 2024a. URL https:// openreview net/ forum? -id--idddid77 W {7HEXg.
Khai Nguyen, Shujian Zhang, Tam Le, and Nhat Ho. Sliced Wasserstein with random-path project- ing directions. In Proceedings of the 41st International Conference on Machine Learning, volume 235, pp. 37879-37899. PMLR, 21-27 Jul 2024b.
Khai Nguyen, Shujian Zhang, Tam Le, and Nhat Ho. Sliced Wasserstein with random-pate project- ing directions. arXiv preprint 240iv:2401. 158888 2024c.
15

---

Published as a conference paper at ICLR 2026
Linh V. Nguyen. Range description for spherical mean transform spaces of constant curvature. Journal d' Analyse Mathématique, 128:191-2144 2016.
Sho Otao and Makoto Yamada. linear time approximation of Wasserstein distance with word embedding selection. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. pp. 15121-15134. Singapore, December 2023. Association for Computa- tional Linguistics. doi:
1002333/1222. mainn.35..........
Minho Park, Taewoong Kang, Jooyeol Yun, Sungwon Hwang, and Jaegul Choo. Spherediff: Tuning- free omnidirectional panoramic image and video generation via spherical latent representation, 2025. URL https: arxiv. org/aba/2504. 14396.
Sung Woo Park and Junseok Kwon. Sphere generative adversarial network based geometric moment matching. In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 4287-4296, 2019.
science. Foundations and Trends@ Machine Learning, 156-6))355-607, 2019. Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data
Michael Quellmalz, Robert Beinert, and Gabriele Steidl. Sliced optimal transport the sphere. Inverse Problems, 39(10)105505,, aug 2023. doi: (10088/136116420//ffffff
Eric Todd Quinto. Support theorems for the spherical radon transform on manifolds. International Mathematics Research Notices, 2006(9):67205 2220006000556555 2006. doi: 10. 155/MMN//2000//72260
Julien Rabin, Gabriel Peyré, Julie Delon, and Marc Bernot. Wasserstein barycenter and its appli- cation to texture mixing. In Alfred M. Bruckstein, Bart M. ter Haar Romeny, Alexander M. Bronstein, and Michael M. Bronstein (eds.), Scale Space and Variational Methods Computer Vision SSVM 2011, Proceedings, volume 6667, pp. 435-446. Springer, 2012. ISBN 978-3-642 24785-9. doi: 10. 10.1007/978-3-6422477777
37.
Mark Rowland, Jiri Hron, Yunhao Tang, Krzysztof Choromanski, Tamas Sarlos, and Adrian Weller. Orthogonal estimation of wasserstein distances. In Proceedings of the Twenty-Scoond Interna- tional Conference on Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pp. 186-195 PMLR, 16-18 Apr 2019.
Meyer Scetbon, Marco Cuturi, and Gabriel Peyré. Low-rank sinkhorn factorization. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 2021.
PMLR, 18-24 Jul
Thibault Séjourné, Gabriel Peyré, and François-Xavier Vialard, Unbalanced optimal transport, from theory to numerics. arXiv preprint arXiv:2211.08775. 2022.
Ashkan Shahbazi, Elaheh Akbari, Darian Salehi, Xinran Liu, Navid NaderiAlizadeh, and Soheil Kolouri. ESPFormer: Doubly-stochastie attention with expected sliced transport plans. In Forty- second International Conference on Machine Learning, 2025. URL pntttss://openi /wpp///eeee.....ee/ net htt:://mmmmmmddduu OMJUBB..
Ondrej Skopek, Octavian-Eugen Ganea, and Gary Bécigneul. Mixed-curvature variational toencoders. In International Conference on Learning Representations, 2020. URL https:
//openrevieww...m////// www.:neet/ooommie--llleeeeee
Li Sun, Zhongbao Zhang, Junda Ye, Hao Peng, Jiawei Zhang, Sen Su, and Philip S Yu. A self- supervised mixed-curvatuu graph neural network. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 4038-4046 2022.
Xiancheng Sun, Senmao Ma. Shengxi Li, Mai Xu, Jingyuan Xia, Lai Jiang, Xin Deng, and Jiali Wang. Spherical-nested diffusion model for panoramic image outpainting. In Forty-second International Conference on Machine Learning, 2025. URL https: attps:://peerevieeee httpp//eeeeee.nettt...eeeeeeeeee forum?i.-..VDFFVf44..
16

---

Published as a conference paper at ICLR 2026
Yuki Takezawa, Ryoma Sato, and Makoto Yamada. Supervised tree- weessssserstein distance. In Ma- rina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pp. 10086-10095, PMLR, 18-24 Jul 2021. html.
URL https:: / http///ooeeddd/ee/e ....... ml press.ppesss..... paress/ve /v http////3///taaaaaalaaaa
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto- encoders. In International Conference on Learning Representations, 2018. URL https:
/ //openreview....
Ba-Hien Tran, Simone Rossi, Dimitrios Milios, Pietro Michiardi, Edwin Bonilla, and Maurizio Filippone. Model selection for bayesian autoencoders. In Advances in Neural Information Pro- cessing Systems, volume 34, pp. 19730-19742, 2021.
Hoang V. Tran, Minh-Khoi Nguyen-Nhat, Huyen Trang Pham, Thanh Chu, Tam Le, and Tan Minh Nguyen. Distance- sasaace-bbased tree-sliced wasserstein distance. In The Thirteenth International Confer- ence on Learning Representations, 2025a. URL https: openrevieeeeee.eeeeeeeeeee net / forum? id= Oi Oigttte tMWwee.
Huy Tran, Yikun Bai, Abihith Kothapalli, Ashkan Shahbazi, Xinran Liu, Rocio Diaz Martin, and Soheil Kolouri. Stereographic spherical sliced wasserstein distances. In Proceedings of the 4 International Conference on Machine Learning (ICML), 2024a. doi: 10.48550/ARXIV24002.
02345. URL https: org/ abs/ 240.. 02345.
Huy Tran, Yikun Bai, Ashkan Shahbazi, John R. Hershey, and Soheil Kolouri. Understanding learning with sliced- wasserstein requires rethinking informative slices, 2024b. URL https:
/ //arxiv.org abs 221..2211...1....11.1.1111111111
Thanh Tran, Hoang V. Tran, Thanh Chu, Huyen Trang Pham, Laurent Ghaoui, Tam Le, and Tan Minh Nguyen. Tree-sliced wasserstein distance with nonlinear projection. In Forty-second International Conference Machine Learning, 2025b. URL https: http:///eeeeeeeeeeeeeeeeeeemeeee / forum? S=S
Viet-Hoang Tran, Trang Pham, Tho Tran, Tam Le, and Tan M. Nguyen. Tree-sliced wasserstein distance on a system of lines, Skii:2400.1372551 [cs.LG], 2024c. URL https://arxiv: org/abs/ 2406. 13725v1. arXiv preprint.
Viet-Hoang Tran, Trang Pham, Tho Tran, Minh Khoi Nguyen Nhat, Thanh Chu, Tam Le, and Tan M Nguyen. Tree-sliced wasserstein distance: geometric perspective. arXiv preprint arXiv:2406. 13725, 2024d.
Viet-Hoang Tran, Thanh T. Chu, Khoi N. M. Nguyen, Trang Pham, Tam Le, and Tan M. Nguyen. Spherical tree-sliced wasserstein distance. In International Conference on Learning Represen- tations (ICLR), 2025c. URL https: lished as a conference paper at ICLR 2025.
Pub-
Viet-Hoang Tran, Thanh Tran, Thanh Chu, Tam Le, and Tan Minh Nguyen. Tree-sliced entropy par- tial transport. In The Thirty-ninth Annual Conference Neural Information Processing Systems, 2025d.
Viet-Hoang Tran, Thanh Tran, Thanh Chu, Duy-Tung Pham, Trung-Khang Tran, Tam Le, and Tan Minh Nguyen. Revisiting tree-sliced wasserstein distance through the lens of the fer- mat-weber problem. In The Fourteenth International Conference on Learning Representations, 2026a. URL https: / /pppeeeeeeieeeeeeeeeeeeeeee net. www.rooorruurooor (et.rummm..kkga00000......
Viet-Hoang Tran, Thanh Tran, Thanh Chu, Duy-Tung Pham, Trung-Khang Tran, Tam Le, and Tan Minh Nguyen. Tree-sliced sobolev IPM. In The Fourleenth International Conference on Learning Representations, 2026b. URL https: hatsss//pppeeeviewe net forummidu HHNOSXaLKF.
Abraham A. Ungar. A Gyrovector Space Approach Hyperbolic Geometry. Synthesis Lectures on Mathematics and Statistics. Morgan & Claypool Publishers, 2008. ISBN 1-111--9-992228 doi: 10.2200/500122ED1 VO1 Y A200005MAS004.
17

---

Published as a conference paper at ICLR 2026
C. Villani. Optimal Transport: Old and New, volume 338. Springer Science & Business Media, 2008.
Minjie Wang, Da Zheng, Zihao Ye, Quan Gan, Mufei Li, Xiang Song, Jinjing Zhou, Chao Ma, Lingfan Yu, Yu Gai, Tianjun Xiao, Tong He, George Karypis, Jinyang Li, and Zheng Zhang. Deep graph library: A graph-centric, highly-performant package for graph neural networks, 2020. URL https: //apss///rxivv org/ abs/ 1909. 01315.
Shen Wang, Xiaokai Wei, Cicero Nogueira Nogueira dos Santos, Zhiguo Wang, Ramesh Nallapati, Andrew Arnold, Bing Xiang, Philip S. Yu, and Isabel F. Cruz. Mixed-curvature multi-relational graph neural network for knowledge graph completion. Proceedings of the Web Conference 2021, WWW '21, pp. 1761-17711 2021.
Yuan Wang, Huazhu Fu, Renuga Kanagavelu, Qingsong Wei, Yong Liu, and Rick Siow Mong Goh. An aggregation-free federated learning for tackling data heterogeneity, 2024. URL https:
wwwrziv.orrg/baee. 2204. 1896222244..2....
Nathaniel S. Woodward, Sang Eon Park, Gaia Grosso, Jeffrey Krupa, and Philip Harris. Product manifold machine learning for physics. In NeurlPS 2024 Workshop on Machine Learning and the Physical Sciences (ML4PS), 2024.
Weijic Xu, Xiaoyu Jiang, Srinivasan Sengamedu Hanumantha Rao, Francis lannacci, and Jinjin Zhao. vONTSS: vMF based semi-supervised neural topic modeling with optimal transport. In Findings of the Association for Computational Linguistics: ACL 2023, pp. 4433-4457, Toronto, Canada, July 2023. Association for Computational Linguisties.
Mingxuan Yi and Song Liu. Sliced wasserstein variational inference. In Proceedings of The 14th Asian Conference on Machine Learning, volume 189 of Proceedings of Machine Learning Re- search, pp.
PMLR, 12-14 Dec 2023.
Liping Yin and Albert Chua. Long range constraints for neural texture synthesis using sliced wasser- stein loss, 2024. URL https: aariiv.iiii.///iii//aaaaaaaaiava org/ wtg/abs////eaaaaceae/..///e...ee abs 1e111.1113..
Youngho Yoon, Inchul Chung, Lin Wang, and Kuk-Jin Yoon. Spheresr: 360deg image super- resolution with arbitrary projection via continuous spherical image representation. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 5677- 5686, June 2022.
Meng Yuan, Yutian Xiao, Wei Chen, Chou Zhao, Deqing Wang, and Fuzhen Zhuang. Hyperbolic diffusion recommender model. Proceedings of the ACM on Web Conference 2025, WWW '25, pp. 1992-2006, 2025.
Wenjie Zheng, Guofeng Zhang, Xiaoran Zhao, Zhikang Feng, Lekang Song, and Huaizhen Kou. Hyperbolic graph wavelet neural network. Tsinghua Science and Technology, 30(4):5511-5555 2025. doi: 010.999//TTT.22440000000 URL https: htwp://c.opp.....cooott. wcw.pen..nnn... www.pppn..co..eoommmmmmmmmme www. camm/arriia cle/
10. 26.99699/// /T0599//TTTTCO 9010032.
18

---

Published as a conference paper at ICLR 2026
TABLE OF NOTATION
R
d-dimensional Euclidean space
1l gd-1
Euclidean norm
(d - -)-dimensional hypersphere
disjoint union
(X)
space of Lebesgue integrable functions on X
P(X)
space of probability measures on X
o()
distribution space
P,V
measures
8(.)
I-dimensional Dirac delta function uniform distribution sample space set of integers from n to m
u(-)
[n.
C(X,y y)
space of continuous maps from X to Y metric in metric space tree metric (rooted) subtree
.(
dr()))
T
tree system
L.
number of Monte Carlo samples
k
number of lines in system of lines or a tree system
R
Radon Transform on Systems of Lines
4--1
(A: 1)-dimensional standard simplex
a
splitting map
T
space of tree systems
N
normal (Gaussian) distribution uniform distribution
u
5
Dirac delta distribution
19

---

Published as a conference paper at ICLR 2026
A
GYROVECTORSPACE
In this section, we summarize some important operations in gyrovectorspaces. Curved Trigonometry.
sinh(tv- -K) V-K
K < 0,
arsinh( (f K
K < 0,
V- K
sing (t) =
t,
K = 0,
sing (t)
t,
= 0, (22)
sin(ty/) K VK
K > 0,
aresin(ty
> 0,
Vk
tanh(t),
K < 0
artanh (t),
K < 0
tank (t)
K
tan K (t) (t)
K = 0
(23)
(t),
K > 0
arctan(t),
K >
Möbius Addition (analogous to vector addition in vector space).
x OOKy (1- (K v)
K (|yll) (1 K K| 2K (ax
K <1|b01/11113
(24)
Scalar Multiplication (analogous vector scalar multiplication in vector space). tank (1 tan (v [V/K] Ir
K 0, 0,
R ||
1 0,
(25)
K + 0, = 0, K 0,
tx,
Gyration.
gyrlz, y) w= 2x 2K v) Ox Ox (y θx v)),
(26)
where Okx== Exponential Map (mapping from tagent space to manifold).
tank (VI)) A A/Kk///llll22))
u θK u,
K + 0, 70,
VIKI 1|
exp (x:)
(27)
K + 0, = 0, K ==0, 0,
u + F.
Logarithm Map (inversion of exponential map).
tank
[ [(-u) 0 1k))
(-u) 1[=
if K # 0, u
/KK22
log K (I)
0,
/ 0, u == if = 0,
u + xr
(28)
Gyrodistance.
V tank
((=))
K 00
dx (x( v)
(29)
2| ul12.
K 0,
Parallel Transport. Along the geodesic from x to y:
K
PTK
(v)
syrly,-rip v,
(30)
A
Geodesics. (analogous lines in vectorspace) The geodesic connecting x and y is parameterized
as
pl() expx (1 log (v)) x Bx & x ((-x) 0x v), LE TR.
(31)
20

---

Published as a conference paper at ICLR 2026
B BASIC PROPERTIES OF RADON TRANSFORM ONTO SEXED-CURVATURE TREE.
We establish the basic properties of Radon transform, including the well-definedness, linearity, and injectivity.
Theorem B.1 (Well-definedness Rg), The operator Rg well-defined
Proof. We compute:
Σ //000000 1/, (1)) dm 2)) do(z) dl
RR /|T
1=1
/ [ ( du(ar, 2)) do(2) dt
Σ I (())) (fo0 8(t 8(t-dd(z,) d)) at) (z)
k
(32)
Σ M
/ ((z)) do)
/ If(=)]
Ea(z.
a(2 Th dz(z)
(33)
S (f(=))
JM
||/l
(34)
co.
(35)
The interchange of integration order in Eq 32 and Eq 33 is justified by Fubini's theorem, as the integrand non-negative and the final integral is finite.
Theorem B.2 (Linearity of R4). The operator R4 is linear, that is,
Ra + bg) = aRq + bR
for all f,g E (M) and a, b R.
Proof. For arbitrary a, b E R and f.g E L' (M), we have:
Ra(a (t,
S (af(z) + bbg())) T), dn(z, 2)) do(z))
= a / (()) T) (()))(2,T)) dm z)) dz(2)
+ b S ((2)a(2, T)x 8( dM(x, ())(()))))) 2)) do(2)
RR f(t, RR r))
(36)
This confirms the linearity of the operator.
Theorem B.3. The operator R is injective.
Proof. We begin by considering the case of single-ray trees. In this setting, the Radon transform reduces to the spherical mean transform in mixed-curvature space. This type of transform has been extensively studied in the context of constant-curvatt 2016). In particular, (Antipov et al., 2011) provide an explicit inversion formula for constant- curvature spaces, which directly establishes the injectivity property in that case. Moreover,
spaces (Agranovsky et al., 2007; Nguyen,
21

---

Published as a conference paper at ICLR 2026
general real-analytic manifold, injectivity follows from the results of (Quinto, 2006), and therefore applies to our setting as a consequence.
We now proceed to the of general trees in mixed-curvature space. Assume that
R 0 A VTET.
(37)
Consider the family of trees with & coinciding rays, denoted by T.. For such trees, the coordinate system and the associated Radon-transformed mass identical across all rays. Define the mapping
g
T
T
rf(t,r") rt.
(38)
immediate that G induces bijective correspondence between T and 7%. Furthermore, we obtain
9 R Rg.
(39)
Consequently, for every Ty Tk.y T,M have
R -.I=GRA--=-
(40)
By the injectivity established in the case of single-ray trees, it follows that f = 0. Hence, we conclude that R is injective.
C METRICITY OF MCSTSW
Theorem C.1 (Metricity of Mixed Curvature Tree Sliced Wasserstein Distance). The Mixed Curva- ture Tree Sliced Wasserstein distance defines metric on (PM)).
Proof. For the Mixed Curvature Tree Sliced Wasserstein distance being a metric, it must compile to three properties:
Positive definitess. As the tree distance a > A v.p E P(T). it follows MCTSW( ()) 2 0. Also, as a v) = A E P(T), MCTSW(U, = 4 P(T). Positive definite follows by the injectivity of Radon transform.
Symmetry. For all v, μ E P(M). we have:
MCTSW(V, S(x We, (RFu, RT
[ W (RG RFv)do(T) TM
= MCTTS((u,)
For H, E P(Sd),, we have Triangle inequality.
P2) + MCTSW(U2,
=f Way (Ro,, RZ) do + JTM
/
War (RTv2, RFvs) do(
[ (War (Rov), R7v22 + War (RFv2, R7vs)) do
Jm
[ War (Rou, RFv3) dT
Jpm MCTSW(U1,
This confirm MCTSW is metric P(M).
22

---

Published as a conference paper at ICLR 2026
D IMPLEMENTATION DETAILS
D.1 COMPUTATION OF DISTANCE FROM A POINT TO A RAY
We first consider the singlc-component case, ie, M CK. If E r.. then (M(2,r3) 0. Otherwise, applying the Generalized Law of Sines (Katok, 1992) gives, for any h E 1 sin(Lzxh) sin( (Zzhrr
sinx (dd(2,,)) = sinx (inx(du(2z
2 sinx (dx((()(z sin sin(<[ (<(xz,
(41)
where the inequality uses 0 < sin(.) < 1.
The first term is computed via curved trigonometry; the angle the second term is computed using left-Möbius addition as an isometry translating (x,, 2, ry) to (0, 2, r8) and preserving angles (Ungar, 2008). Thus,
(3, 2}
cos(<(ry,
(42)
||
The distance is then computed with
dm(2, r²) = sin (sinx(d (iinssn()))))s((((())) sin(<(zz,r)))
(43)
For a product space assuming (y x))² + O.
d = dek, (2, +
Edeg,
deg, (2,2)²
(44)
D.2 CLOSED-FORM OF WASSERSTEIN DISTANCE ON MIXED-CURVATUM TREE
In this subsection, we derive the closed-form of (R5,, Assume
m(x) Σ V2 8(z -aj)
and p(z)
Σ
Mg 8(z ag)
(45)
This representation unify the support of discrete distributions, allowing zero mass. Consider the mixed curvature tree T T px......'' Let 6j = dm(x,,,) 02) Vj E [1, n). and let co = 0. Without loss the generality, we assume that
0 co S c1 S 02 S S Cn
(46)
For 0 < j S n and 1 < S k, consider all points zj = (cz, ) on the mixed curvature
(i)
tree T. We have To Fo
(1) (2)
= (k)
(i) is the intersection between the ball and x
B, {z d2( dr) = oj) and the ray res forall S j S n. We compute
R.q r())) for E [0, +00) and sissk
Rov((, x2) = /
u(z) T) 6(1
do(z)
M
=
Ty 8(2 a¹) a(z a(z T), ott du 2)) dz
vj I
Σ
((2,T), 8(2 - az) [( dm(ar, 2)) do(z).
(47)
1=1
Hence,
R9v(1,12) = -0 (0
if t & (ci....,c. c.
{a T);
(48)
if = cj for some
23

---

Published as a conference paper at ICLR 2026
The formula for R is achieved similarly.
For j S n and S i k, let
8 a(aj. T)
1 - alaz, T) Mj.
(49)
Consider T as graph with nodes i (the root). The edges are
zj (1)))
= 1 for 0 S j < n, sis k, where 10 - r for all
isisk, 1 13 n,
with length E(e)5) = Cj cj-1. For j > 0, the subtree [ax))) contains all nodes z/ with 2 j.
The Tree-Sliced Wasserstein distance on T is
Wa.(Ruu RGH) = Σ Σ
(g -1-1
v(T(af))) ((((o))))
(50)
==
From
v(((x)))) = Σ
(
(T((=)))) (0) = Σ 4
p
p=j
we obtain
Wa, (RFURR91
Σ
Co C9-1)
Σ 1 (1
(51)
i=1 j=1
p=j
c c--1)
Σ a(pp. T) (Up Hp)
(52)
E EXPERIMENTAL DETAILS
E. GRADIENT FL.WW
Our gradient flow experiment aims at matching distributions on the product manifold Eg x P2 x DD- We use mixture of 6 WNDs, each with Σ 0.011, centered distance 0.5 from the origin and equally spaced (equiangular) between consecutive centers, as the target distribution v. The source distribution y WND centered at the origin. We use full-batch Riemannian gradient descent (Boumal, 2023) to minimize MCTSW(,, /4). Both the source and target consist of 2,400 data points. We tune the learning rate for each method while keeping the total number of projections fixed at 1,800.
The baselines are:
SWamhieet SW distance in ambient space R
Prod-TSW: 1-product metric using the MCTSW variant on each component space
Prod-SW: 1-product metric using the geodesic SW in each component
E.2 VARIATIONAL. AUTO-ENCODER
Baselines. The vanilla VAE (Kingma & Welling, 2013) uses Euclidean space as latent space, with the objective consisting of Binary Cross Entropy as reconstruction loss and KL-divergence as reg- ularizer. SWAE (Kolouri et al., 201 8a) replaces the KL--diergencc with Sliced- Waseed-tessrrteinn in Eu- clidean space. S-VAE (Davidson et al., 2018) uses the sphere as latent space and KL-divergence as regularizer. STSW-VAE replaces the KL-divergence on the sphere with the Sphere Tree-Sliced Wasserstein distance (Tran et al., 2025c). H-VAE uses hyperbolic space as latent space and KL- divergence regularizer. HSW- VAE replaces the KL-divergence with the hyperbolic sliced Wasser- stein distance (Bonet et al., 2022a). M-VAE uses MCS as latent space and KL-divergence as regu- larizer. The latent spaces for the Fuclidean, Spherical, Hyperholic, and MCS cases are Eg, Sg, Pg.
24

---

Published as a conference paper at ICLR 2026
and E3 D3 P3. respectively. All methods use a mixture of 10 WNDs (or vMF for the spherical case, and Normal distribution for the Euclidean case) as the prior distribution. For the MCS and Hyperbolic cases, after passing through the encoder, embeddings are mapped to the manifold us- ing the exponential map. With this parameterization, we can utilize Euclidean optimizers since the parameters lie in Euclidean space.
Dataset. The CIFAR-10 dataset (Krizhevsky & Hinton, 2009) is widely used benchmark in ma- chine learning and computer vision. It consists of 60,000 color images, each of size 32 x 32 pixels, divided evenly into 10 distinct classes. Each class contains 6,000 images representing common ob- jects including airplanes, cars, birds, cats, deer, dogs, frogs, horses, ships, and trucks. The dataset split into 50,000 training images and 10,000 test images, enabling clear evaluation and benchmark- ing of image classification models. CIFAR-10 10 is known for moderate size and diversity, making it suitable for rapid experimentation with machine learning algorithms such as convolutional neural networks (CNNs) and other classification methods. It was originally developed by researchers the subset of the larger Tiny Images dataset, with manual labeling to ensure
University of Toronto quality, Its balanced classes and diversity make it challenging but accessible dataset for developing and testing image recognition models.
Training, We adopt the training protocol from (Tran et al., 2025c), using the Adam optimizer (Kingma & Ba, 2014) with learning rate set 0.001. The total number of projections is fixed at 2,000 for all SW-based regularizers, and for TSW-based methods, we set the number of trees to 200.
Architecture for MCTSW-VAE.
Encoder:
I E IR 3x32x32 = Conv2dg2 - ReLU
Conv2d32 ReLU
- Conv2de1 - ReLU
Conv2de4 ReLU
Conv2d12s - ReLU
Conv2d12s Flatten
FC512 - ReLU - FCg
exp_map M
Decoder:
M log_map FC512
FC2048
ReLU
E.3 GRAPH SELF-SUPERVISED LEARNING
Dataset. Cora (Kipf & Welling, 2017) is a benchmark citation network where nodes are research papers and edges represent citation links. Each paper is represented by a binary bag-of-words vector (1 if dictionary term appears, 0 otherwise), with labels corresponding to subject areas. We utilize the preprocessed version DGL (Wang et al., 2020).
Hyperparameter settings. We use the Adam optimizer for all methods with learning rate 5 x 10-4. The balancing hyperparameter is selected from (0.01, 0.05, 0.1, 0.5, 1} for best performance. The number of estimators for random forest is set to 200. Number of trees in each TS method 500 with rays on a tree.
Architecture. We follow the same architecture as in (Liu et al., 2024) to use GCN consist of 2 graph convolutional layer with hidden dimension of 256.
E.4 HARDWARE SETTINGS
All experiments were conducted single NVIDIA A100 (80GB) GPU server.
E.5 RUNTIME ANALYSIS
The training times for our three experiments involving variants of Wasserstein distance are summa- rized in Table 4, 5, and 6. It is worth noting that, in the gradient-flow setting (Table 4), our imple-
25

---

Published as a conference paper at ICLR 2026
Table 4: Training time and performance for the gradient flow experiment.
Method
SW __ambiiet
Prod-TSW
Prod-SW
MCSTSW
Time (s)
158
617
468
303
log W2 (1)
0.33
0.34
0.31
-3.65
Table 5: Training time and accuracy for the graph self-supervised learning task.
Method
S-TSW-SSGE
H-TSW-SSGE
E-TSW-SSGE
MC-TSW-SSGE
Time (s)
63.15 + 0.60
63.36 1.08
63.06 = 0.95
62.93 + 0.84
Accuracy (1)
78.3 + 0.2
75.1 + 0.2
77.9 = 0.3
79.9 = 0.5
mentations of Prod-TSW and Prod-SW apply TSW/SW independently to each subspace, whereas MCTSW exploits a highly parallelizable tree-based construction. This yields substantially lower runtime for MCTSW despite operating in a mixed-curvature space. The same parallelization strat- egy explains why the MCTSW-based variants in Tables 5 and Table achieve training times that are competitive with other latent-space baselines while achive superior performance.
E.6 ABLATION STUDY
We provide ablation studies over these hyperparameters on graph self-supervised learning task in Tables and 2. In notation, mack Ko > 0 and m mo mean that we use ma subspaces whose curvatures are in Ku. -Ko + 2Ko - Ko + 1Ko ..... Ky _ which we implement in code as aranee (-m0/ 22, m0/2) 2 KC / m0. Note that mazk = 0 recovers
the Euclidean space (E-TSW-SSGE in Table 2). The results in Table 7 and Table 8 suggest that moderate values of both the curvature magnitude and the number of subspaces are most beneficial.
For the parameter & (number rays per tree), additionally conduct an ablation study on the graph self-supervised learning task with mazk = 4 and m 32. The results in Table 9 show that the performance stable across a wide range of
F
DESIGN CHOICES
F.1 PROJECTION FUNCTION AND SPLITTING MAP
We implement three additional variants, covering different combinations of
projection rules (including spherical projection (our version) and orthogonal projection, and
mist-iiittinution strategies across rays (by distance to the orthogonal projection (our ver- sion) or to the spherical projection),
along with a detailed comparison of the time complexity for each method. Across all variants, the two fundamental tree factors-roo position and ray directions incorporated carefully to ensure that no information is lost in projection. The performance and running time of these variants on the Gradient Flow task with the same setting as in Section 5.1 are shown in Table 10. is evident that our version of MCTSW outperforms the three others and has lower run time.
Looking at the runtime, line, resulting in time complexity of O( log n) instead of O(In log n), as with the spherical projection our design choice. Furthermore, the spherical splitting functionn, while having the same time complexity O(Ldmnk), has higher hidden constant, leading to significantly higher run time. Note that in our notation, L is the number of trees used for Monte Carlo estimation, & denotes the number of rays per tree, is the number of support points, d denotes the dimension of each component, and m is the number of components. In our experiments, we use 360, d = 2, 3, 2400, and 5.
noteworthy that orthogonal projection requires a different sort for each
26

---

Published as a conference paper at ICLR 2026
Table 6: Training time per epoch and test BCE for the VAE task.
Method
SWAE
STSW-VAE
HSW-VAE
MCSTSW-VAE
Time (slepoch)
9.64 0.09
7.25 0.01
9.76 + 0.18
10.62 0.05
Test BCE (1)
0.6043 + 0.0005
0.6026 = 0.0009
0.6012 + 0.0006
0.6000 + 0.0002
Table 7: Ablation study on the maximum magnitude of curvature K when m 32.
mArK
4
Accuracy (1) 77.9 + 0.3 77.7 0.1 78.3 + 0.2 79.9 + 0.5 77.3 0.5
Moreover, would like to note that the metricity properties can be verified similarly to the proof in Appendix and Appendix for the two versions using spherical projection, while the metricity of the two variants using orthogonal projection requires further investigation.
F.2 TREE SIMPLIFICATION
We developed version of MCTSW that sample trees without our simplification specified in Sec- tion and conducted an additional experiment, comparing the two version. The only difference between the two lies in how the splitting function is computed. In our version, there exists closed- form expression for the distance from a point a geodesic ray (detailed in Appendix D.I), whereas the non-simplified version has to approximate this distance numerically.
For this case, we adopt ternary search algorithm (CP Algorithm) to approximate (2, min 2 minge-- dM(22, & ((-z) @ y),) and we terminate the search after 10 iterations.
The results the gradient-flow experiment are reported in Table 11. The two versions achieve similar performance in terms log W2, while the non-simplified variant significantly slower in wall-clock time. This empirically supports choice of the simplified formulation.
F.3 TREE TOPOLOGY
we have further evaluated the impact of the tree topology by implementing an MCTSW variant in the non-simplified setting with a sequence (chain-like) like) topology, i.e., a sequence of lines joined at the intersection of every two consecutive lines. The trees are sampled by drawing the intersection points from a wrapped normal distribution centered at the origin.
As this variant consumes substantially more memory, for a fair comparison between the two topologies, namely T. 50, d 2, 3, n 100, and k 5. The results in Table 12 show that the star topology used in main experiments clearly outperforms the chain-like topology in terms of log W2.
smaller gradient-floy flow configuration
27

---

Published as a conference paper at ICLR 2026
Table 8: Ablation study on the number of subspaces m when mark =
m
32
128
Accuracy (1) 76.7 + 1.1 79.9 + 0.5 78.7 + 0.2
Table Ablation study on the number lines per tree l when mark and m
k
4
16
32
128
256
512
Accuracy (1)
79.6 = 0.2
79.9 0.5
79.5 = 0.3
79.6 + 0.4
79.8 + 0.1
79.7 = 0.4
79.8 0.6
79.5 + 0.3
Table 10: Results on gradient-flow task log W2 and run time.
Combination
log W2 (1)
Run time (s)
Orthogonal split + Orthogonal projection
-1.7517
383
Orthogonal split Spherical projection (ours)
-3.6576
303
Spherical split + Orthogonal projection
-2.6038
1350
Spherical split Spherical projection
-3.5546
1110
Table 11: Comparison of simplified and non-simplified MCTSW (gradient flow task).
Method
log W2 (L)
Run time (s)
Simplified
33.6576
303
Non-simplified
-3.6569
3479
Table 12: Comparison of MCTSW with star topology and chain-like topology.
Tree topology
log Wa (L)
Star
-3.3788
Chain-like
-1.7140
28


Reviews:
[{"review_id": "PlkDj1kImP", "reviewer": "Reviewer_HAxH", "summary": "The central hypothesis is that the Tree-Sliced Wasserstein framework can be generalized to mixed curvature spaces (MCS). The core technical contributions include a mixed-curvature tree system $\\mathcal{T}$ and a Radon Transform $\\mathcal{R}^{\\alpha}$ for MCS, which maps mass from the manifold $\\mathcal{M}$ onto $\\mathcal{T}$. The authors claim this construction yields a closed-form, computationally efficient solution for the 1-Wasserstein distance on the tree. They propose that this MC-TSW metric will more faithfully capture the joint geometric structure of MCS distributions than separable baselines (like product-space SW/TSW), leading to superior performance.", "questions": "- The success in Table 1 suggests the joint projection (via $d_{\\mathcal{M}}$ and $\\alpha$) is doing all the work.  Does this imply the tree structure itself is less important than the projection function?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is nicely written and tackles a well-motivated and non-trivial problem. The authors correctly identifies the limitations of naive baselines. That is, ambient-space metrics ignore the geometry, while separable \"product\" metrics (like Prod-TSW) ignore joint structure across components. This work appears to be the first to build a TSW-like distance specifically for these joint spaces.\n\n\n- Good experimental results. For example, \n\n  - the gradient flow experiment provides strong evidence for the paper’s central claim. In Table 1, the MC-TSW method successfully converges to a low $\\log W_2$ error ($-3.65$). In stark contrast, the separable baselines, Prod-SW and Prod-TSW, fail entirely (positive $\\log W_2$ errors). This result seems to suggest that the target distribution have cross-component correlations that only a joint metric like MC-TSW can capture.\n\n\n  - the VAE and graph SSL experiments (Tables 2 & 3) further show that combining the proposed MC-TSW regularizer with an MCS latent space yields SOTA or highly competitive results, outperforming both KL-based regularizers and single-curvature (Euclidean, spherical, hyperbolic) models. This validates both the new metric and the underlying geometric hypothesis.\n\n- The algorithm has the same asymptotic cost as TSW, parallelizes well on GPUs.", "weaknesses": "- The Radon transform (def 3.6) projects every point $z \\in \\mathcal{M}$ to the same coordinate $t = d_{\\mathcal{M}}(x, z)$ on all $k$ rays of the tree. The only thing distinguishing the projection of $z$ onto different rays is the splitting map $\\alpha(z, \\mathcal{T})_i$. This spherical projection collapses all angular and component-specific directional information of $z$ relative to the root $x$, reducing it to a single scalar distance. This appears to be a geometrically lossy projection, and its trade-off is not discussed.\n\n- Moreover, the paper's construction is limited to star-shaped trees at a single point $x$. The implementation in algo 1 simplifies this further by sampling rays that are axis-aligned, i.e., each ray extends in exactly one component $j$ of the product space. While the projection onto this tree (via the joint distance $d_{\\mathcal{M}}$ and splitting map $\\alpha$) is still joint, the tree structure itself is highly constrained and separable. This axis-aligned, star-shaped probe may not be the most effective geometry for capturing complex, non-axis-aligned correlations, even if it is computationally convenient.\n\n- The $O(L n \\log n)$ component of the complexity is a direct result of this information-collapsing projection. Because all points $a_j$ are projected to the *same* coordinate $c_j = d_{\\mathcal{M}}(x, a_j)$ on every ray, a single sorting of the $n$ values $\\{c_j\\}$ is sufficient to compute the $W_1$ on the tree (Eq. 21). This tractability is thus achieved by collapsing all per-ray projection information into a single scalar distance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AhhopgY4SL", "reviewer": "Reviewer_NuB5", "summary": "By utilizing mixed-curvature spaces that enable better data representations aligned with complex underlying structures—rather than relying solely on the Euclidean setting—this paper introduces the Mixed-Curvature Tree-Sliced Wasserstein, a computationally efficient discrepancy measure between distributions. The proposed technique builds upon the ideas behind the classical Sliced-Wasserstein (SW) framework, leading the authors to introduce a Radon-type transform. The paper demonstrates the effectiveness of this new tool through experiments involving generative flow models and variational autoencoders.", "questions": "- Although I understand the general idea, I got lost and could not clearly identify where the stereographic projection $\\rho_k$ is utilized in the methodology. Measures are assumed already in the hyperspher $F$ Poncare ball P\n- Could the authors clarify what they mean by “consistency of operators across geometries” in line 155?\n- Is there any intuition or guideline on how to choose the parameters $k$ and $m$?\n\nStyle Suggestions:\n- The paper alternates between using MC-TSW and MCTSW. I suggest unifying the notation throughout the text.\n- I would reserve the symbol $\\rho$ exclusively for the projection in Equation (3). For instance, I recommend using a different symbol for the Sliced-Wasserstein operator in Equation (8), as well as later in Equation (18) and similar instances.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is very well organized.\n- The paper is thorough and complete: it presents solid theoretical results, clearly introduces the necessary notions and new definitions, establishes key properties, and provides simplifications that enhance its applicability. It also includes extensive references to related work, provides time complexity analyses, and validates the proposed framework through different experiments.\n- The idea is novel and has the potential to be applied to data analysis across various disciplines.", "weaknesses": "The only potential weakness is that the paper might appear overly technical at times; for example, it would benefit from including a brief explanation of how to obtain the closed-form expression in Equation (21).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KYaPFkAA0O", "reviewer": "Reviewer_So7Z", "summary": "This paper proposes an algorithm to construct a tree when the data lies in a mixed-curvature space. Then, using this tree metric, this paper proposed computing the tree-wasserstein distance efficiently.", "questions": "See the weakness section.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* This paper proposed a novel method to construct the tree.\n* The proposed methods can capture the structure of Muxture-curvature spaces, and the effectiveness of the proposed methods was demonstrated in several tasks: gradient flow and self-supervised learning.", "weaknesses": "* This paper did not evaluate the approximation error of the proposed method and the original Wasserstein distance. For instance, can the proposed method more accurately approximate the Wasserstein distance than the tree-sliced Wasserstein distance with the tree metric constructed by the clustering method [4]? This paper evaluated the performance of the proposed method in several tasks; however, these results are influenced by various factors. It would be good if this paper could demonstrate the effectiveness of the proposed method in the simplest task.\n* Various methods have been proposed for constructing tree metrics for the tree-Wasserstein distance, e.g., [1,2,3,4], but this paper did not discuss the relationship between the proposed method and these existing approaches. Although the proposed method appears to differ from these existing approaches, it is necessary to discuss its relationship to them.\n* In the experiments, what methods did this paper use for TSW to construct the tree metric?\n* This paper did not evaluate the computational efficiency of the proposed method. In lines 378-388, this paper discusses the time complexity; however, can this paper also report the time required for training in the experiments?\n\n## Reference\n[1] Indyk et. al., Fast image retrieval via embeddings. In International Workshop on Statistical and Computational Theories of Vision 2003\n\n[2] Takezawa et. al., Supervised tree-wasserstein distance. In International Conference on Machine Learning 2021.\n\n[3] Lin et. al., Tree-wasserstein distance for high dimensional data with a latent feature hierarchy, In International Conference on Learning Representation 2021.\n\n[4] Le et. al., Tree-sliced variants of wasserstein distances, In Neural Information Processing Systems 2019", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)