Tree-sliced Sobolev IPM
Authors: Viet-Hoang Tran
Thanh Tran
Thanh Chu
Duy-Tung Pham
Trung-Khang Tran
Tam Le
Tan Minh Nguyen
Venue: ICLR 2026 Poster (2026)
Keywords: sliced optimal transport
sobolev ipm
tree-sliced wasserstein distance
tree wasserstein distance
Primary area: learning on graphs and other geometries & topologies
TL;DR: We introduce TS-Sobolev, a tree-sliced metric derived from a closed-form regularized Sobolev IPM, broadening the practical use of Wasserstein distance on tree-metric spaces.

Abstract: Recent work shows Tree-Sliced Optimal Transport to be an efficient and more expressive alternative to Sliced Wasserstein (SW), improving downstream performance. Tree-sliced metrics compare probability distributions by projecting measures onto tree metric spaces; a central example is the Tree-Sliced Wasserstein (TSW) distance, which applies the $1$-Wasserstein metric after projection. However, computing tree-based $p$-Wasserstein for general $p$ is costly, largely confining practical use to $p=1$. This restriction is a significant bottleneck, as higher-order metrics ($p > 1$) are preferred in gradient-based learning for their more favorable optimization landscapes. In this work, we revisit Sobolev integral probability metrics (IPM) on trees to obtain a practical generalization of TSW. Building on the insight that a suitably regularized Sobolev IPM admits a closed-form expression, we introduce TS-Sobolev, a tree-sliced metric that aggregates regularized Sobolev IPMs over random tree systems and remains tractable for all $p \ge 1$; for $p>1$, TS-Sobolev has the same computational complexity as TSW at $p=1$. Notably, at $p=1$ it recovers TSW exactly. Consequently, TS-Sobolev serves as a drop-in replacement for TSW in practical applications, with an additional flexibility in changing $p$. Furthermore, we extend this framework to define a corresponding metric for probability measures on hyperspheres. Experiments on Euclidean and spherical datasets show that TS-Sobolev and its spherical variant improve downstream performance in gradient flows, self-supervised learning, generative modeling, and text topic modeling over recent SW and TSW variants. Our code is available at [https://github.com/thanhquangtran/TS-Sobolev](https://github.com/thanhquangtran/TS-Sobolev).

Published as a conference paper at ICLR 2026
TREE-SLICED SOBOLEV IPM
Viet-Hoang Tran'
Thanh Q. Tran"
Department of Mathematics National University of Singapore hoang. tranviet@u @u. nus edu
Department of Computer Science National University of Singapore thanhtran@u nus. edu
Thanh Chu
Duy-Tung Pham
Department of Computer Science National University of Singapore thanh. chuêu. nus. edu
FPT Software Al Center Vietnam temgpa10@fpt.co
Trung-Khang Tran
Tam Le
Department of Computer Science National University of Singapore TTK01 O 06 nus.edu
Department of Advanced Data Science Institute of Statistical Mathematics tam i sm. ac ip
Tan Minh Nguyen Department of Mathematics National University of Singapore tarmn& nus edu.sg
ABSTRACT
Recent work shows Tree-Sliced Optimal Transport be an efficient and more expressive alternative to Sliced Wasserstein (SW), improving downstream perfor- mance. Tree-sliced metrics compare probability distributions by projecting mea- sures onto tree metric spaces; central example is the Tree-Sliced Wasserstein (TSW) distance, which applies the I-Wasserstein metric after projection. How- ever, computing tree-based Wassersteii for general is costly, largely confining practical use to p This restriction significant bottleneck, as higher-order metrics (p 1) preferred in gradient-based learning for their more favorable optimization landscapes. In this work, we revisit Sobolev integral probability met- ries (IPM) trees to obtain practical generalization of TSW. Building the insight that a suitably regularized Sobolev IPM admits a closed- -orsed-form expression, we introduce TS-Sobolev, tree-sliced metric that aggregates regularized Sobolev IPMs over random tree systems and remains tractable for all 2 1; for > 1, TS Sobolev has the same computational complexity as TSW at p 1. Notably, at 1 it recovers TSW exactly. Consequently, TS-Sobolev serves as a drop- replacement for TSW in practical applications, with an additional flexibilityy in changing Furthermore, we extend this framework to define correspond- ing metric for probability measures on hyperspheres. Experiments on Euclidean and spherical datasets show that TS-Sobolev and its spherical variant improve downstream performance gradient flows, self-supervised learning, generative modeling, and text topic modeling recent SW and TSW variants. Our code is available at https: //tpt/gttbh//////////////h com/ wwwnhhanngtaannnnnnanauaaaaaan Sobolev..
INTRODUCTION
Comparing probability measures a foundational problem numerous scientific fields where data often represented as distributions. For example, documents in natural language processing
Equal contribution
Co-last author. Please correspond hoang w...eee..uu....iiii..... and tammn @ mus..du.ssg

---

Published as a conference paper at ICLR 2026
be treated distributions of words or topics (Sparck Jones, 1972; Kusner al., 2015; Yurochkin et al., 2019), and 3D shapes in computer vision are commonly modeled as point clouds, which discrete distributions of data points (Achlioptas et al., 2018; Hua et al., 2018; Wu et al., 2019). Optimal Transport (OT) has emerged as a powerful framework for this purpose (Villani, 2008; Peyné et al., 2019), as it defines a metric between distributions that inherently respects their underlying geometry. This key advantage has driven its broad adoption in fields like machine learning (Bunne al., 2022; Fan al., 2022; Takezawa et al., 2022), data valuation (Just et al., 2023; Kessler al., 2025), multimodal data analysis (Park et al.. 2024; Luong et al., 2024), statistics (Mena & Niles- Weed, 2019; Wang et al., 2022; Liu et al., 2022; Nietert et al., 2022), and computer vision (Lavenant et al., 2018; Saleh et al., 2022). However, a significant drawback of Optimal Transport (OT) is its computational complexity. For discrete measures supported by n samples, standard algorithms scale as On log n) (Peyré et al., 2019), rendering OT impractical for large datasets.
Sliced Optimal Transport The substantial computational burden of Optimal Transport (OT) led to the development of the Sliced Wasserstein (SW) distance as a powerful and efficient approxima- tion (Rabin et al., 2011; Bonneel et al., 2015; Nguyen, 2025). At its core, SW simplifies the problem by leveraging the fosed-form solution of me-dimensional OT. It projects high-dimensional proba- bility measures onto random one-dimensional subspaces, computes the simple transport cost in each "slice", and averages these costs. This procedure reduces the computational complexity that of sorting, (n log n)) (Peyré et al., 2019), while faithfully preserving important statistical and topo- logical properties of the original metric (Nadjahi et al.. 2020; Bayraktar & Guo. 2021: Goldfeld & Greenewald, 2021). The success of this paradigm has inspired a broad family of extensions, includ- ing those using structured projections (Kolouri et al., 2019; Deshpande et al., 2019; Nguyen al., 2020: Ohana et al., 2023; Nguyen et al., 2023) and adaptations for non-Euch non-Euclidean geometries such as spheres (Bonet et al., 2022; Quellmalz et al., 2023) and hyperbolic space (Bonet et al., 2023b).
Tree-Sliced Optimal Transport. A key limitation of classical SW, however, is that one- mne-dimeniional projections can be insufficient for capturing complex geometric structures inherent in high- dimensional data. This has spurred research into alternative slicing domains beyond simple lines, with explorations across various metric spaces like Fuclidean subspaces (Alvarez-Melis et al., 2018; Paty & Cuturi, 2019; Niles-Weed & Rigollet, 2022), graphs (Le et al., 2022), and non-Euclidean manifolds (Tran et al., 2024c; Bonet et al., 2023a; Lin et al., 2025). Among these, methods based on tree metrics have emerged as a particularly effective approach. The Tree-Sliced Wasserstein (TSW) distance, introduced by Tran et al. (2024d), capitalizes on closed-form OT solutions on tree metric spaces (Le et al., 2019; Indyk & Thaper, 2003; Indyk, 2001). TSW thus achieves a compelling bal- ance: it maintains the computational tractability of SW while better representing the structure of the data. Subsequent works have improved TSW (Tran et al., 2025a;c; 2026b) and extended it to other domains, such as mixed-curvature spaces (Pham et al., 2026) and spheres (Tran et al., 2025d).
Sobolev Integral Probability Metric. The computational efficiency of the Tree-Sliced Wasser- stein (TSW) distance is direct result of the closed-form solution for the 1-Wasserstein distance on tree metric spaces (Le et al., 2019). A significant limitation of this framework, however, is that this analytical solution does not extend to orders p > 1, which restricts the applicability of TSW. This restriction is problematic because higher-order metrics are often preferred in gradient-based learning tasks. Specifically, p-Wasserstein with p > offers strict convexity (Santambrogio, 2015; Villani, 2003) and smoother gradients compared to the p = 1 case, properties that are known to fa- cilitate more stable and efficient optimization (Peyré et al., 2019). To address this gap, we leverage the framework of Integral Probability Metrics (IPMs), a powerful class of distances for comparing probability measures (Müller, 1997). IPMs function by finding a critic function from a predefined class that maximally discriminates between two distributions, a versatile principle with numerous applications in machine learning and statistics (Sriperumbudur et al., 2009; Gretton et al., 2012: Liang, 2019; Uppal et al., 2019; 2020; Nadjahi et al., 2020; Kolouri et al., 2020). A theoretically important instance is the Sobolev IPM, which constrains the critic function to a unit ball defined by the Sobolev norm (Adams& & Fournier, 2003). This specific metric has been instrumental in theo- retical analyses, such as studying convergence rates and error bounds in generative models (Liang, 2017; 2021; Singh et al., 2018). Despite its theoretical appeal, the standard Sobolev IPM lacks closed-form expression, which has historically limited its practical use. A recent breakthrough by Le et al. (2025) overcomes this challenge by introducing a regularized Sobolev IPM for probability
2

---

Published as a conference paper at ICLR 2026
measures supported on tree. This novel formulation yields closed-form solution that is both computationally efficient and valid for any order 2 1.
Contributions. Motivated by the expanding Tree-Sliced Wasserstein (TSW) framework and recent advances in Sobolev IPMs on tree metric spaces (Le et al., 2025), this paper introduces a novel family of tree-sliced distances. Our primary contribution is scalable generalization of TSW for probability measures in Euclidean spaces and on the sphere. By leveraging the closed-form solution of the regularized Sobolev IPM, our proposed distance is efficiently computable for any order p 2 1, overcoming key limitation of previous TSW variants and leveraging the optimization advantages associated with higher-order metrics. The the paper is organized as follows:
1. In Section 2, tree metric spaces and detail the closed-form solution for the regularized Sobolev IPM on trees. We also revisit the Tree Systems framework and the Radon transform that underpin the Tree-Sliced Wasserstein (TSW) distance.
we establish the building blocks for our method. We review the theory of
2. In Section 3, we introduce the Tree-Sliced Sobolev IPM (TS-Sobolev) for measures on Eu- clidean spaces and the Spherical Tree-Sliced Sobolev IPM the sphere. We establish their metric properties, prove they provide a scalable generaliza- tion of TSW for any order p 1, and analyze their computational complexity.
for measures on
3. Section 4 presents experiments on Euclidean and spherical data that validate our method's practical effectiveness and efficiency, followed by our conclusion in Section 5.
Supplementary materials, including detailed theoretical background, full proofs, and extended perimental results (setups, tables, and figures), are available in the Appendix.
BACKGROUND ON SOBOLEV INTEGRAL PROBABILITY METRIC AND TREE-SLICIN
This section covers the two foundational concepts behind our method. We first review the Sobolev Integral Probability Metric (Sobolev IPM) and its efficient closed-form solution on tree metric spaces. We then describe the tree-slicing framework that projects measures from Euclidean space onto these tree metric spaces, thereby enabling the use of the efficient Sobolev IPM.
2.1 SOBOLEV INTEGRAL PROBABILITY METRIC
Tree Metric Spaces. A free metric space (T, dr) is a continuous space built from a tree T= (V,E) with vertices V C Rd and edges E. Each edge is assigned a non-negative length we. Crucially, the space T includes not only the vertices V but also every point along the edges E. Its tree metric, de- noted dr, is the unique path distance between any two points on the tree (Semple & Steel, 2003b; Le al., 2019). The unique path between points æ and y is denoted [r. y/. This structure gives rise to canon- ical Borel length measure, w, where the measure of any path equals its length: (([,, v)) dy(r: v). Fi- nally, the subtree rooted at x, denoted A A(x), the set of all points y whose path from the root must pass through æ, i.c., A(z) {y E T : E [r. v]}.
A(wa)
A(z1)
66
E
#5
0g
010
62
64
700
08
61
23
44
178
211
Figure 1: An illustration of tree metric space. The tree, rooted rt consists of nodes 2; and edges e; with weights w.- A probabil- ity distribution on the tree then assigns mass to nodes. The subtree A(r) is the collection of all points lying along the edges in the sub- tree rooted at z. For example, A(r) includes the entire tree, A(x1) ) (red) includes all points edges e2 63, and ca, and A (x²g) (blue) in- cludes all points on edges e e11, and e12-
Comparing Measures on Trees. A central task is to define a distance between probability measures on a tree T. Let P(T) denote the collection of all prob- ability measures on T (i.c., those with a total mass of 1). For any two measures μa, v E P(T), we write P(u, u) for the set of all valid couplings æ between them. A popular framework for defining
3

---

Published as a conference paper at ICLR 2026
distance is the p- Wasserstein distance:
( inf xEP(HM) TxT
J+xT dr(z, dz(z.))
Wp(u,v)
(1)
While powerful, computing this distance is generally expensive. However, for the special case of = 1, the distance on tree admits a fast, closed-form solution (Le et al., 2019):
WI(,,D) v) = lu(A(z)) - v(A())) w(d)).
(2)
Crucially, for 1, no similar closed-fone solution is known to exist. This computational bottle- neck for higher-order Wasserstein distances on trees motivates our exploration of the Sobolev IPM, which provides tractable alternative for all p 2 1.
Sobolev Integral Probability Metric. The Sobolev Integral Probability Metric (IPM) is a distance between measures on tree T, defined using concepts from functional analysis. The metric is situated within the Sobolev space WTL.((T,,) .w),, which consists of functions that have well-defined tree derivatives. This space is equipped with a Sobolev norm, which in turn defines a unit ball of functions, denoted B(p'). For rigorous treatment of these concepts, we refer to Appendix B. The Sobolev IPM is formally defined by finding the maximum discrepancy between two measures, J and D, as evaluated by critic function f constrained to this unit ball:
(1, [(x)
f ((v) u(dy)
Spl,,, v) = sup SEB(y')
((d))))))))))))))))))()))) mldx)
where y' is the conjugate of p (where y' E [1, oo] satisfies 1/p + 1/p/ = 1; 1, then x0).
Regularized Sobolev Integral Probability Metric. The computational intractability of the standard form in Equation (3) motivates the use of a Regularized Sobolev Integral Probability Metric (IPM), which admits direct closed-form solution (Le et al., 2025). For continuous measures over the tree and for any order S < oo, this solution is given by:
D)(, J, ((<=)>>)u((()))) Iu(A(z)) v(A(x))) w(dr),
(4)
where ú(x) = 1 + (A(T)) is a weight function.
For the practical case of discrete measures supported on the tree's nodes V. this integral simplifies to an efficiently computable sum over the tree's edges E:
Sp D)P
Σ cEE
- -((7))
(5)
where % = (y E T: S [r. y/} is the set of points whose path from the root contains edge e. The coefficient B, is pre-computable scalar, making the overall computation efficient:
We
log 1 1 w(TT) + w(Ye) + w.
if = 2,
Be =
(6)
P (1 + ((2)))--P otherwise.
2-P
Analyzing Equation (4) shows that > 1 yields gradients that scale with the eтror magnitude [p(A(z)) - (A(())), unlike the constant gradients associated with p 1, thereby facilitating smoother optimization consistent with established properties of Wasserstein (Santambrogio, 2015; Peyré al., 2019). Furthermore, the term vi'x' down- weights global gradients near the root, al- lowing the optimization to focus refining fine-grained local structures at the leaves. This weight- ing constitutes unique optimization advantage TTSSooolev over standard y- Wasserssteinn metrics. We provide comprehensive analysis of these optimization advantages in Appendix F.8.
2.2 THE TREE-SLICING FRAMEWORK
While the Regularized Sobolev IPM defined in Section 2.1 offers a powerful metric for measures supported trees, data in machine learning typically resides in Fuclidean space Rdd, To bridge this
4

---

Published as a conference paper at ICLR 2026
gap, utilize the Tree-Slicing framework (Tran et al., 2024d). This framework provides mecha- nism to project probability measures from R onto continuous tree metric spaces, thereby allowing us to leverage the closed-form solutions of trre-brrsese bes-d metrics for high-dimensional Euclidean data.
Tree System. A line in R an clement of Id x gd-1, and a system of l: lines is an element of (Rd x 8d-1)k We denote system of lines by - {4 [ where each line I is defined by source point 2; E R and direction vector 0g S6-1, with parameterization I; L.; for t E IR
free system is a system of lines endowed with an additional tree structure. To highlight this structure, write T {4) The ground set of this system, denoted T, is the set of all points on all lines in T, formally defined as T {(x, 4,) E IR x T This ground set forms a continuous tree metric space as defined Section 2.1. The space of tree systems sharing common tree structure is denoted by Ty (or simply T), equipped with probability distribution a induced by a random sampling procedure over the lines.
+ tz0, for some to R).
Radon Transform Tree Systems. Let C(R. x TV Ak-1) be the set of continuous splitring maps from R2 T to the (k 1)-simplex 1 For f E L
(IR),, define R.y f: R by
RS R9,(2,4) = { a(y; 6(x (y - Ii, 0:)) dy.
The operator
Re: f (R9f) TET!
is called the Radon Transform Tree Systems. This operator is injective.
Tree Sliced Wasserstein distance. The Tree-Sliced Wasserstein (TSW) distance is metric between probability measures, defined the expected I-Wasserstein distance between the measures after projection onto tree system. Recent variants include the Tree-Sliced Wasserstein Distance on Systems of Lines (TSW-SL) (Tran al., 2024d) and its generalization, the Distance-based Tree- Sliced Wasserstein Distance (Db-TSW) (Tran et al., 2025a), which we refer to collectively as TSW in this paper. For two probability measures v E P(R4) with corresponding densities Ju. Su. the TSW distance is formally defined as:
TSW(H, J W1 W1 fR (T)...
TSW is a valid metric on PIR)) and can be efficiently approximated via Monte Carlo sampling, thanks to the closed-form solution for the W-Wasserstein distance on trees Equation (2).
It is crucial to distinguish the variants from earlier work that also uses the TSW name (Le et al., 2019; Sato al., 2020; Yamada al., 2022; Takezawa et al., 2022). Those methods were primarily designed for static-support measures where the data points are fixed. In contrast, the TSW-SL and Db-TSW formulations are specifically built for dynamic-supper measures, where the points change during optimization. Our goal is to leverage this powerful tree-slicing framework to develop a new Sobolev IPM-based metric suitable for these dynamic tasks that mitigates the computational bottleneck of the p- Wassersseiin distance trees.
TREE-SLICED SOBOLEV IPM
In this section, we propose the Tree-Sliced Sobolev IPM (TS-Sobolev) framework for probability distributions on Euclidean spaces and the sphere. We establish its theoretical guarantees, clarify its connections to existing metrics, and characterize its computational complexity.
3.1 TREE-SLICHE SOBOLEV IPM
For probability measures μ. E P(RRd) with respective densities fn. and fe, the Tree-Sliced Sobolev IPM defined the average regularized Sobolev IPM between i and induced by tree-metric projections. Given tree system T e T and a splitting map E C(Rd x Tr At-1). the Radon transform RR maps fu. fo densities R2 fu, R4fu f on T, inducing measures HT, UT E P(T). We then evaluate the regularized Sobolev IPM (HT, 27) as in Equation (5) and define the resulting distance as the expectation of this quantity over with respect to the sampling distribution
5

---

Published as a conference paper at ICLR 2026
Definition 3.1 (Tree-Sliced Sobolev IPM). The Pree-Sliced Sobolev IPM of order p [1, o0),
denoted as TS-Soboleyp. between μ. V/ E P(IR)) ) is defined by:
T5- = Sp(UT, uT) do do(T))
(9)
3.2 PROPERTIES OF TREE-SLICED SOBOLEV IPM
Metricity of TS-Sobolevy- We recall the Euclidean group E(d) and state that Fidaainvariance of the splitting map ensures that the Tree-Sliced Sobolev IPM defines a metric. Let IR be equipped with the Euclidean norm || |l2. The Euclideun group E(d) is the group of all distance-preservingg
transformations of R4; it is the semidirect product T(d) x O(d) the translation group T(d) {Q Rdxd; Q'Q La}- Every
{o - + g E E(d) can be written as = (Q. with Q E O(d) and v E R.. acting on E R by gy Qy v. This action extends the space of tree systems T by gT {gl, ) 1=1 for lines represented as (I; fs)) we set gl, == (Qx; UU
E R$) and the orthogonal group ((d)
U, Q+:), which preserves the underlying tree structure.
splitting map a C(R2 Tr. AR-1) 11) is E(d)-iivariant
a(gı, g7) = a T) for all IR" T E TEY, and g E E(d).
(10)
Invariance under the Euclidean group, E(d), is desirable property for distances between proba- bility measures on Rd, Standard metrics like the 2- 2-Wasserstein and Sliced p- p-Wasserstein distances possess this H( B(d)-invariance. However, for the Tree-Sliced Sobolev IPM, this property is even more fundamental: guarantees that TS-Sobolev a valid metric on ((R.)..
Theorem 3.2. TS- Sobolev is an B(d)-invariant metric on P(RR)).
The proof for Theorem 3.2 is presented in Appendix D.1.
Connections to TSW. The Tree-Sliced Sobolev IPM serves as a natural generalization of the Tree- Sliced Wasserstein (TSW) distance. Notably, for the order p = the TS-Sobolev IPM recovers the TSW distance exactly. For any order p E [1, oo), it is upper-bounded by the TSW distance.
Theorem 3.3. For any μt, V E P(Rd)) and E oo): (S-Sobolee,((,1p)) S TSW(u, v), where equality holds for p 1, Le.. TS- TS-Sobolev v) = (s,(,,
The proof for Theorem 3.3 appear in Appendix D.2.
Remark 3.4. We note that the values of TS-Sobolev, and TSW depends on the choice of tree system (T), sampling distribution (o)), and splitting map (a.). In this paper, utilize the specific choices established in the Distance-basee based TSW (Db-TSW) framework (Tran et al., 2025a). Therefore, the TSW we analyze is precisely Db-TSW. To maintain readability, these dependencies suppressed in the notation for T--Soboley,, with detailed description of the framework available in Appendix C.
Computation of Tree-Sliced Sobolev IPM. The intractable integral in Equation (9) approxi- mated using Monte Carlo estimate:
$
1 L Σ
TS (Soboley,(,,))
Sp(UTo uT:)
(11)
where we sample L tree systems (Ti from a distribution a. Let μr and y be discrete measures with and m support points, respectively, and assume >> m. The overall computational complex- ity is (Lkn log + Lkdn), where is the number of lines per tree and d is the data dimension. This is identical to the complexity of TSW variants like Db-TSW (Tran et al., 2025a), as the extra step of computing the coefficients 3r per Equation (6) adds a negligible ((Lkn) cost.
A key advantage is that this complexity holds for any order p [1, co). resolving the computational intractability of higher-order TSW. Empirically, the runtime of TS Sobolev is nearly identical to that of the first-order DS-TSW, confirming its efficiency. A runtime analysis is provided in Appendix F.1.
The practical application of TS-Sobolev depends on its hyperparameters: the number of trees L, lines per tree k, and the order p. Prior work (Tran et al., 2025a) shows that using multiple lines

---

Published as a conference paper at ICLR 2026
(A: > 1) is crucial for capturing complex data topology. This creates natural trade-off between increasing for expressiveness and increasing L to improve the precision of the Monte Carlo esti- mate by reducing its variance. The convergence rate of this estimate with respect to L is formalized in Theorem 3.5. A detailed sensitivity analysis for the tree parameters L and l is presented in Appendix F.6, while the influence of the order p is analyzed in Appendix F.7.
Theorem 3.5. The approximation error of T'SSobolev decreases a rate of O
We defer the proof for Theorem 3.5 to Appendix D.3.
3.3 EXTENSION TO THE SPHERICAL SETTING
The TS-Soboley, framework extends to measures on the hypersphere, μ.v E P(S"), by using spherical tree systems (Tran et al., 2025d). We provide a brief derivation of the resulting metric below, deferring complete treatment to Appendix E. The core idea to use the spherical Radon transform on spherical free systems to map the densities of and onto a given tree T, which induces the projected measures jrt and VT-
Definition 3.6 (Spherical Tree-Sliced Sobolev IPM). The Spherical Tree-Sliced Sobolev IPM of order p E [1, se), denoted as STS-Soboleyp- between μl, V E P(3d) is defined by
1
(T = (L, S uT) da(T))
(12)
detailed derivation of STS-Soboley,. along with a full analysis of its properties, provided in Appendices E.2 and E.3, respectively. Notably, for the p 1 case, STS- STS-Sobolev, recovers the Spherical Tree-Sliced Wasserstein (STSW) distance exactly, as our implementation adopts the spe- eific splitting map and tree sampling methodology from Tran et al. (2025d).
EXPERIMENTAL RESULTS
In this section, we empirically evaluate our proposed methods across a diverse range of applications to demonstrate their effectiveness in both Euclidean and spherical settings. For the Euclidean setting, we conduct experiments on gradient flows and generative modeling with diffusion models. For the spherical setting, our evaluation focuses on peff-supervised learning benchmark. Additionally, we assess our methods on a topic modeling task, for which we provide results in both domains.
4.1 EVALUATION ON EUCLIDEAN DATA
Gradient Flow on Rd. In this experiment, we apply our methods to gradient flow task, which seeks to find a path of distributions μh that minimizes a distance 'D between an initial source (4) = N (0, T) and fixed target The evolution of this path is governed by the update rule 0.p = - v),, where μt is the distribution at time and VD(u,, v)) is the corresponding distance gradient. Our evaluation is conducted
the 8 Gaussians and Gaussian 30d datasets, where we benchmark our TS- TS-Sobolev variants (p E {1.2, 1.5, 2}) against a comprehensive suite of baselines. These include Siiced-Wasserstinn (SW) methods- -uchh as vanilla SW (Bonneel et al., 2015). MaxSW (Deshpande et al., 2019). LCVSW (Nguyen & Ho, 2023), and SWGG (Mahey et al., 2023) as well as recent Tree-Sliced (TSW) distances like TSW-SL (Tran al., 2024d), Db-TSW, and Db-TSW- (Tran al., 2025a). We assess performance by measuring the Wasserstein distance to the target at intervals up to 2500 iterations, with detailed results available in Table 1.
The results demonstrate that our TS-Sobolev, methods (p E {1.2, 1.5, 2}) achieve better conver- gence compared to baselines. While some SW variants perform well initially, our methods consis- tently improve and ultimately outperform their Wasserstein--asse TSW counterparts, For example, on the Gaussians datasct, TS-Sobolev achieves a final distance of 8.88 x 10-7 outperform- ing the strongest baseline. Db-TSW (2.50 x 10-5). Similarly, the Gaussian 30d dataset, TS- Sobolev1.2 has the best final distance of 1.40, surpassing both TSW-SL (1.93) and Db-TSW (1.78).
Diffusion Models. This experiment applies our proposed TS-Sobolev distance to the task of train- ing Denoising Diffusion Generative Adversarial Networks (DDGANs) (Xiao et al., 2021) for un- conditional image synthesis. Following the approach of Nguyen et al. (2024), integrate our
7

---

Published as a conference paper at ICLR 2026
Table 1: Average Wasserstein distance (multiplied by 10-1 for Gaussian 30d) between source and larget distributions of 10 runs. All methods use 100 projecting directions.
Gaussians
Gaussian 30d
Iteration
Iteration
Methods
500
1000
1500
2000
2500
500
1000
1500
2000
2500
3.97e-2
6.48e-3
1.08e-3
1.09e-3
1.08e-3
2.93
2.87
2.80
2.72
2.64
MaxSW
4.66e-2
3.53e-2
2.74c-2
2.33e-2
2.08e-2
2.24
2.53
2.68
2.68
2.64
SWGG
7.57e-3
7.00e-5
5,80e-5
5,68e-5
5.71e-5
2.72
2.74
2.74
2.74
2.74
LCVSW
7.50e-4
5.42e-4
5.53e-4
5.58e-4
5.43e-4
2.85
2.71
2.58
2.45
2.33
TSW-SL
1.92c-2
7.42c-4
1.34c-6
1.33c-6
1.17c66
2.48
2.31
2.16
2.04
1.93
Db-TSW
8.18e-5
2.51e-6
2.26c-6
2.24c-6
2.50c-6
2.44
2.24
2.07
1.90
1.78
T 2.12e-2
1.95c-6
1.25e-6
1.08e-6
8.88c-7
2.38
2.10
1.85
1.62
1.40
TS T
2.93e-2
1.17c-3
3,28c-6
2.27c-6
2.03e-6
2.44
2.25
2.02
1.77
1.51
TS- TS-Sobolevz 3.05e-2
8.69c-3
1.43e-4
4.54c-6
3.50c-6
3.21
3.49
3.58
3.61
3.68
A comparison of DDGAN mod- els on the CIFAR-10 unconditional genera- tion benchmark, showing Fréchet Inception Distance (FID) scores and per-epoch training times averaged over 10 runs.
Table 2:
Table 3: Accuracy of the linear classifier
coded (E) features and projected (P) features on $
Method
Acc. E(%) 1
Acc. P( ((%)
Hypersphere
79.76
74.57
Model
FID
Time/Fjoch (S)
SimCLR
79.69
72.78
DIGINN
al., 2021)
3.64
SSW
70,46
64.52
SW-DD (Ngoyen et 2024)
2.90
74
DDW-DD (Nguyen 2024)
S3W
78.54
73.84
EBSW-DD (Nguyen al., 2024)
2.87
76
RI-S3W (5)
79.97
74.27
RDSW-DD (Nguyen 2024)
2.82
IWRPSW.DD
2.70
77
ARI-S3W (5)
79.92
75.07
DD (Tran et al., 2024d)
2.83
80
D6-TSW-DD
2025a)
STSW
80.53
76.78
D6-TSWDDD (Trun 2025a)
2.53
85
STS-Sobolevi.6
79.88
76.07
TS-Soholes, ,DD (ours)
2.302 + 0.004
84
STS-Soholeva
80.6
77.65
SS-Sobolevg-DD (ours)
distance into the Augmented Generalized Mini-batch Energy (AGME) loss function. We bench- mark TS-Sobolev against several Sliced and Tree-Sliced Wasserstein-based DDGAN variants, with full experimental details available in Appendix F.3.
The results, summarized Table 2, show that our methods yield notable improvements in sample quality. Notably, both TS-Sobolev 1.-DD and TS-Sobolev,-DD surpass the strongest baseline, Db- TSW-DD (Tran et al., 2025a), reducing the Fréchet Inception Distance (FID) by 0.228 and 0,253, respectively. This gain in sample quality is achieved without a trade-off in efficiency, as our methods have comparable training times to other Tree-Sliced variants.
We attribute TS-Sobolev's performance gains to its ability to preserve fine-graie mne-grained image details, as theoretically and empirically demonstrated in Appendix F.8.
4.2 EVALUATION ON SPHERICAL DATA
Self-Supervised Learning (SSL). Previous work by (Wang & Isola, 2020) demonstrated that the contrastive objective can be separated into two key components: an aligmment loss, which ensures that embeddings of similar inputs remain close, and uniformity loss, which prevents collapse by encouraging representations distribute more evenly. Building on the idea of (Bonet et al., 2022), substitute the Gaussian kernel used in the uniformity term with our proposed method.
C
- 1
(STS-Soboley, STS-Soboley,
Alignment loss
Uniformity loss
8

---

Published as a conference paper at ICLR 2026
Table 4: Log of the Wasserstein distance between source and target distributions over 10 runs on mixture of 12 vMFs.
Epoch
Methods
50
100
150
200
SSW
-2.439 + 0.053
-2.787
0.040
-2.909 + 0.041
-2.979 + 0.037
-3.014 + 0.034
S3W
-2.022 = 0.036
-2.211 + 0.045
-2.284 - 0.056
-2.290
-2.289 1 0.064
RI-S3W (1)
-2.094 0 0.028
-2.488 + 0.028
-2.693 =
2.814
-2.900 0.026
RI-S3W (5)
-2.433 = 0.029
-2.790 + 0.023
-2.999 =
0.022
-3.093 0.021
ARI-S3W (30)
-2.612 = 0.043
-2,942 + 0.029
-3.090 0.035
+
-3,270 0.047
LSSOT
-2.078 + 0.030
-2.444 0.023
-2.546 + 0.023
-2.582 0.023
-2.598 + 0.021
STSW
-2.693 0.030
-3.171
0.041
3.376
-3,488 + 0.049
-3,549 + 0.072
STS-Sobolevi.s
-3.099 0.032
-3,324 0.050
-.427=
-3,499 0.064
-3,540 = 0.078
STS-Soboleya
-3.081 = 0.026
-3,376 + 0.058
-3.513 0.094
-3.578 0.108
-3.616 + 0.123
where y. = U(S represents the uniform distribution on the unit sphere S 2 ,=,2 E (ax(d+1) denote the embeddings of two augmented views the same sample and 0 serves as a weight to balance the alignment and uniformity terms. Following the approach in Bonet et al. (2022); Tran al. (2024c; 2025d), we apply this objective to pretrain ResNet18 He et al. (2016) encoder on CIFAR-10 10 Krizhevsky et al. (2009) for 200 epochs. After pretraining, a linear classifier is trained on top of the frozen encoder to evaluate learned features.
As shown in Table 3, our proposed STS- variants demonstrate superior performance com- pared to both tree-sliced and standard sliced baselines. STS-Sobolev, achieves the best overall accuracy (80.6% Encoded / 77.65% Projected), outperforming its direct tree-based counterpart, STSW (Tran et al., 2025d). Furthermore, our method significantly improves upon standard spherical slicing approaches, such as SSW (Bonet al., 2022) and S3W variants (Tran et al., 2024c).
Gradient Flow on the sphere. In this task, our objective is to learn the target distribution v from a source distribution j by minimizing d(v, μ.) where d is the distance metric such as SSW (Bonet et al., 2022), S3W (Tran et al., 2024c), LSSOT (Liu et al., 2025) and STSW (Tran et al., 2025d). Consistent with prior works Bonet al. (2022); Tran al. (2024c; 2025d), Fisher distributions (vMFs) with 2400 samples as the target distribution. Optimization is carried using projected gradient descent Bonet al. (2022)on the sphere with full-batch training. We report in Table 4 the log 2asssersttiin distance at epochs 50, 100, 150, 200, and 250, averaged over 10 runs. Across all epochs, our proposed STS-Sobolev consistently outperforms the baselines. In particular, while STS STS-Sobolev 1.5 achieves the best result at epoch 50, STS-Soboleyy demonstrates the strongest overall performance at other epochs. At the final epoch, STS-Soboleva achieves a distance of 3.666 + 0.123, surpassing its tree-sliced counterpart STSW (-3.549 + 0.072) as well as the sliced baselines LSSOT 22599 = 0.021) and ARI-S3W (-3.270 = 0.047).
mixture of 12 von Mises-
Topic Modeling. Topic modeling task (Blei al., 2003) seeks to automatically extract distinct themes from collections of text documents, revealing the underlying structure of a corpus. Recent neural approaches typically employ variational autoencoder (VAE) setup, in which the optimiza- tion balances accurate document reconstruction with a regularization that encourages the inferred topic distributions to resemble a chosen prior (Srivastava & Sutton, 2017). Inspired by Nan et al. (2019); Adhya & Sanyal (2025), we propose replacing the conventional KL-divergence regularizer with Wasserstein-based alternative. This leads to the following objective:
inf By, [CE(x, x)] + ATS SSTS--oboley, (0). P2.0
where CE represents the cross-entropy between the input document (in bag-of-words represen- tation) and its reconstruction x. The variational posterior (,/x)) is generated by encoder yp, and the decoder & maps topic mixtures / back to word distributions to form &. The encoder yp can be adapted output 0 in either IR or gd which allows for direct evaluation of our methods in both the Euclidean and spherical settings.
We evaluate proposed TS-Sobolev: and STS-Sobolev; topic models against comprehensive set of basclines, with all results summarized in Table 5. In the Euclidean setting, our model is benchmarked against several modern Sliced Wasserstein (SW) and Tree-Sliced Wasserstein (TSW)

---

Published as a conference paper at ICLR 2026
Table Average topic coherence CV(() BBC and M10 over 10 runs.
Method
BBC
M10
LDA (Blei et al., 2003)
0..574 0.68840.018
0..414
ProdLDA (Srivastava Sutton, 2017)
0.49140.017
WTM (Nan et al., 2019)
0.74140.034
0.403
Euclidean setting
SW (Bonneel al., 2015)
0.760
0.484
RPSW (Nguyen et al., 2024)
0.77510.022
0.472
EBRPSW (Nguyen et al., 2024)
0.777 #0.019
0.490
TSW-SL (Tran et al., 2024d)
0.796
0.45640.000
Db-TSW (Tran et al., 2025a)
0.787 10.041
0.458
TS-Soboleva (Ours)
0.805
0. 0.497
Spherical setting
SSW (Adhya & Sanyal, 2025)
0.75540.056
0.40840.018
S3W (Tran et al., 2024c)
0.700
0.102
LSSOT (Liu et al., 2025)
0.743+0.051
0.395+0.049
STSW (Tran et al., 2025d)
0.752+0.031
0.373+0.013
STS-Sobolev; (Ours)
0.776
0.423
variants from recent works (Bonneel al., 2015; Nguyen et al., 2024; Tran al., 2024d; 2025a). In the spherical setting. we compare against recent spherical slicing methods, specifically SSW (Adhya Sanyal, 2025), STSW (Tran al., 2025d), S3W (Tran et al., 2024c), and LSSOT (Liu et al., 2025). To provide broader context, we also include results from three foundational topic models: LDA (Blei et al., 2003), ProdLDA (Srivastava Sutton, 2017), and WTM (Nan et al., 2019).
To assess the quality of the discovered topics, 2015), where higher values indicate better topic quality. The results, summarized in Table 5, show that proposed methods consistently achieve the highest topic coherence scores in all settings. In the Euclidean setting, our TS-Sobolev model obtains a top score of 0.805 on the BBC dataset, surpassing the best baseline score of 0.796. On the M10 dataset, it also leads with score of 0.497 compared to the baseline best of 0.490. This strong performance extends to the spherical setting, where STS-Sobolev, model attains scores of 0. 776 BBC and 0.423 M10, outperforming the strongest respective basclines (0.755 and 0.408). Full details on the experimental setup are provided in Appendix F.5.
the standard coherence metric Cv (Röder et al.,
CONCLUSION
This paper introduced the Tree-Sliced Sobolev Integral Probability Metric (TS-Sobolev) and its spherical variant, S7S-Sobolev, novel approaches for comparing probability measures. Our work generalizes the Tree-Sliced Wasserstein (TSW) framework by leveraging regularized Sobolev IPM, enabling the efficient computation of tree-sliced distances for any order p 2 1. We presented a for- mal derivation of these metrics and provided comprehensive theoretical guarantees, including proofs of metricity and the formal connection the original TSW distance. Experimental evaluations show that TS-Sobolev and -To--Sooleev consistently outperform state-of-the-an Sliced Wasserstein and Tree-Sliced Wasserstein methods across various tasks, including topic modeling and training dif- fusion models. Crucially, these performance gains are achieved with no additional computational overhead, as our methods maintain a runtime comparable to existing TSW techniques. A limitation of our current work is that TS Sobolee is not designed to compare unbalanced measures, where input measures may have different total masses. Therefore, a promising future direction is to design an extension of our method for the unbalanced optimal transport setting.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
We thank the area chairs and anonymous reviewers for their comments. TL gratefully acknowledges the support of the JST-BOOST program (FY2025), JSPS KAKENHI Grant number 23K11243, and Mitsui Knowledge Industry Co., Ltd. grant. TT acknowledges support from the Application Driven Mathematics Program funded and organized by the Vingroup Innovation Fund and VinBigData.
This research project is supported by the National Research Foundation Singapore under the AI Singapore Programme (AISG Award No: AISGG-TTC-2023-012--5
This research project is supported by the Ministry of Education, Singapore, under the Academic Research Fund Tier (FY2023) (A-8002040-00-000 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 our experiments are provided in the supplementary materials of the paper. The details of our experimental settings and computational infrastructure are given in Appendix F. All datasets that we used 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
Panos Achlioptas, Olga Diamanti, Joannis Mitliagkas, and Leonidas Guibas. Learning representa- tions and generative models for 3D point clouds. In International conference on machine learning, pp. 49. PMLR, 2018.
Robert A Adams and John JF Fournier. Sobolev spaces, volume 140. Elsevier, 2003.
Suman Adhya and Debarshi Kumar Sanyal. S2WTM: Spherical for topic modeling. In Proceedings of the 63rd Anmual Meeting of the Association for Compu- tational Linguistics (Volume Long Papers), pp. 23211 1-2225.... Association for Computational Linguistics, July 2025. doi: 10. 10.18655/11/2225.acllle 1131.
autoencoder
David Alvarez-Melis, Tommi Jaakkola, and Stefanie Jegelka. Structured optimal transport. In International Conference on Artificial Intelligence and Statistics, pp. 1771-1780. PMLR, 2018.
Erhan Bayraktar and Gaoyue Guo. Strong equivalence between metrics of Wasserstein type. Elec- tronic Communications in Probability, 26(none):1 13, 2021. doi: 10.1214/11-ECP383
David M Blei, Andrew Y Ng. and Michael I Jordan. Latent dirichlet allocation. Journal of Machine Learning Research, 3:993-10222 1022, Jan 2003.
Clément Bonet, Paul Berg, Nicolas Courty, François Septier, Lucas Drumetz, and Minh-Tan Pham. arXiv preprint arXiv:2206.0888 08780, 2022.
Spherical
Clément Bonet, Lactitia Chapel, Lucas Drumetz, and Nicolas Courty. Hyperbolic SiicedWasseesstee via geodesic and horospherical projections. In Topological, Algebraic and Geometric Learning Workshops 2023, pp. 370-370. PMLR, 2023a.
Clément Bonet, Laetitia Chapel, Lucas Drumetz, and Nicolas Courty. Hyperbolic sliced-Wasserse Wasserstein via geodesic and horospherical projections. In Proceedings of 2nd Annual Workshop Topology, Algebra, and Geometry in Machine Learning (TAG-ML), pp. 334-370. PMLR, 2023b.
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.
Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause. and Marco Cuturi. Proximal optimal trans- port modeling of population dynamics. In International Conference on Artificial Intelligence and Statisties, pp. 1-1-28... PMLR, 2022.
11

---

Published as a conference paper at ICLR 2026
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International confer- ence on machine learning. pp. 1990-2999. PMLR, 2016.
Ishan Deshpande, Mhan-Ting Hu, Ruoyu Sun, Ayis Pyrros, Nasir Siddiqui, Sanmi Koyejo, Zhizhen Zhao, David Forsyth, and Alexander Schwing. Max-sliced Wasserstein distance and its use for GANs. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10648-10656, 2019.
Jiaojiao Fan, Isabel Haasler, Johan Karlsson, and Yongxin Chen. On the complexity of the optimal transport problem with graph-structured cost. In International Conference on Artificial Intelli- gence and Statistics, pp. 9147-9165. PMLR, 2022.
Ziv Goldfeld and Kristjan Greenewald. Sliced mutual information: A scalable measure of statistical dependence. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 17567-17578 Curran Associates, Inc., 2021.
Derek Greene and Pádraig Cunningham. Practical solutions to the problem of diagonal dom- inance in kernel document clustering. In Proceedings of the 23rd International Conference on Machine Learning. ICML '06, pp. 377-384, New York, NY, USA, 2006. Association for Computing Machinery. ISBN 1595933832. doi: 1444441148844.14488888 URL https:
//doi. 1141/11441111 1143892.
Arthur Gretton, Karsten M Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel (wo-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.
Allen Hatcher. Algebraic topology. Cambridge University Press, 2005.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Binh-Son Hua, Minh-Khi Tran, and Sai-Kit Yeung. Pointwise convolutional neural networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 984-993, 2018.
Piotr Indyk. Algorithmic applications of low-distortion geometric embeddings. In Proceedings 42nd IEEE Symposium on Foundations of Computer Science (FOCS), pp. 10-33, 2001.
Piotr Indyk and Nitin Thaper. Fast image retrieval via embeddings. In International workshop on statistical and computational theories of vision, volume 2, pp. 5, 2003.
Hoang Anh Just, Feiyang Kang, Tianhao Wang. Yi Zeng, Myeongseob Ko, Ming Jin, and Ruoxi Jia.
LAVA: Data valuation without pre-specified learning algorithms. In The Eleventh International Conference on Learning Representations, 2023.
Samuel Kessler, Tam Le, and Vu Nguyen. SAVA: Scalable learning data valuation. In The Thirteenth International Conference Learning Representations, 2025.
Kinga, Jimmy Ba Adam, al. A method for stochastic optimization. International conference on learning representations (ICLR), volume 5, pp. 6. San Diego, California; 2015.
Soheil Kolouri, Kimia Nadjahi, Umut Simsekli, Roland Badeau, and Gustavo Rohde, Generalized sliced Wasserstein distances. Advances neural information processing systems, 32, 2019.
Soheil Kolouri, Nicholas Ketz, Andrea Soltoggio, and Praveen K. Pilly. Sliced Cramer synap-
tic consolidation for preserving deeply learned representations. In International Conference Learning Representations, 2020.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images.
Technical Report, University of Toronto, 2009.
12

---

Published as a conference paper at ICLR 2026
Matt Kusner, Yu Sun, Nicholas Kolkin, and Kilian Weinberger. From word embeddings to document distances. In International conference on machine learning, pp. 957-966 2015.
Hugo Lavenant, Sebastian Claici, Edward Chien, and Justin Solomon. Dynamical optimal transport on discrete surfaces. In SIGGRAPH Asia 2018 Technical Papers, pp. 250. ACM, 2018.
Tam Le, Makoto Yamada, Kenji Fukumizu, and Marco Cuturi. Tree-sliced variants of wasserstein distances. Advances in neural information processing systems, 32, 2019.
Tam Le, Truyen Nguyen, Dinh Phung, and Viet Anh Nguyen. Sobolev transport: A scalable metric for probability measures with graph metrics. In International Conference on Artificial Intelligence and Statistics, pp. 9868. PMLR, 2022.
Tam Le, Truyen Nguyen and Kenji Fukumizu. Scalable unbalanced Sobolev transport for measures on graph. In International Conference on Artificial Intelligence and Statistics, pp. 8521-8560.
PMLR, 2023.
Tam Le, Truyen Nguyen, and Kenji Fukumizu. Generalized sobolev transport for probability mea- sures on a graph. In Forty-first International Conference on Machine Learning, 2024.
Tam Le, Truyen Nguyen, Hideitsu Hino, and Kenji Fukumizu. Scalable sobolev IPM for probability measures on a graph. In Forty-second International Conference on Machine Learning, 2025. URLhttps: /tpp///peeeeeeeeeeeeeeeee net forum? boaummid=Vhhhn vi-
Tengyuan Liang. How well can generative adversarial networks learn densities: A nonparametric view. arXiv preprint, 2017.
Tengyuan Liang. Estimating certain integral probability metric (IPM) as hard as estimating under the IPM. arXiv preprint, 2019.
Tengyuan Liang. How well generative adversarial networks learn distributions. Journal of Machine Learning Research, 22(228):1-41 2021.
Ya-Wei Eileen Lin, Ronald R. Coifman, Gal Mishne, and Ronen Talmon. tance for high dimensional data with latent feature hierarchy. In The Thirteenth International Conference on Learning Representutions, 2025.
dis-
Lang Liu, Soumik Pal, and Zaid Harchaoui. Entropy regularized optimal transport independence criterion. In International Conference on Artificial Intelligence and Statistics, pp. 1 11247-11279. PMLR, 2022.
Xinran Liu, Yikun Bai, Rocio Diaz Martin, Kaiwen Shi, Ashkan Shahbazi, Bennett Allan Landman, Catie Chang, and Soheil Kolouri. Linear spherical sliced optimal transport: A fast metric for com- paring spherical data. In The Thirteenth International Conference on Learning Representations, 2025. URL WURLLttsss///peeeeeeeee /ttp://pppeeeeeeeeeeeeeee http://wwwwee..tttueeeeee festt/formmmdddoeeeee Pine:d==fuu
ZAxywx.
Manh Luong, Khai Nguyen, Nhat Ho, Reza Haf, Dinh Phung, and Lizhen Qu. Revisiting deep audio-text retrieval through the lens of transportation. arXiv preprint 2rXiv:2405.1000 2024.
Guillaume Mahey, Lactitia Chapel, Gilles Gasso, Clément Bonet, and Nicolas Courty. Fast optimal transport through sliced generalized wasserstein geodesics. In Thirty-sevnnth Conference Neu- ral Information Processing Systems, 2023. URL https: /ttp://////pneeeeee.. net / forum? idd= N3XuYdvhNW.
Gonzalo Mena and Jonathan Weed. Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem. In Advances in Neural Information Processing Systems, pp. 4541 4541-4551, 2019.
Alfred Müller. Integral probability metrics and their generating classes of functions. Advances in Applied Probability, 29(2):429-444 1997.
Kimia Nadjahi, Alain Durmus, Lénaïc Chizat, Soheil Kolouri, Shahin Shahrampour, and Umut Simsekli. Statistical and topological properties of sliced probability divergences. Advances in Neural Information Processing Systems, 33:20802-20812, 2020.
13

---

Published as a conference paper at ICLR 2026
Feng Nan, Ran Ding, Ramesh Nallapati, and Bing Xiang. Topic modeling with Wasserstein autoencoders. In Proceedings of the 57th Anmual Meeting of the Association for Computa- tional Linguisties, pp. 2345-6381. Association for Computational Linguisties, July 2019. doi:
1///655/1//1191666 URL https:: / /acl anthology org/p1 9- 1600/.
Khai Nguyen. An introduction to sliced optimal transport, 2025. URL https: /awpiv.xx.ii.////../... abs/ /2008 12519.
Khai Nguyen and Nhat Ho. Sliced Wasserstein estimation with control variates. arXiv preprint 0000:23305.00402,, 2023.
Khai Nguyen, Nhat Ho, Tung Pham, and Hung Bui. Distributional sliced Wassersennn and applica- tions to generative modeling. arXiv preprint arXiv:2002.074: 073667 2020.
Khai Nguyen, Tongzheng Ren, Huy Nguyen, Litu Rout, Tan Minh Nguyen, and Nhat Ho. Hi- erarchical sliced Wasserstein distance. The Eleventh International Conference on Learning Representations, 2023.
Khai Nguyen, Shujian Zhang, Tam Le, and Nhat Ho. Sliced Wasserstein with random -att project- ing directions. arXiv preprint arXiv:2401.15888 2024.
Sloan Nietert, Ziv Goldfeld, and Rachel Cummings. Outlier-robust optimal transport: Duality, struc- ture, and statistical analysis. In International Conference on Artificial Intelligence and Statistics, pp. 1169 1-11719. PMLR, 2022.
Jonathan Weed and Philippe Rigollet. Estimation of Wasserstein distances in the spiked trans- port model. Bernoulli, 2663)))666662888 2022.
Ruben Ohana, Kimia Nadjahi, Alain Rakotomamonjy, and Liva Ralaivola. Shedding PAC- Bayesian light on adaptive sliced-Wassereee Wasserstein distances. In Proceedings of the 40th International Conference Machine Learning, 2023.
Shirui Pan, Jia Wu, Xingquan Zhu, Chengqi Zhang, and Yang Wang. Tri-party deep network repre- sentation. In Proceedings of the Twenty-Fifth International Joint Conference on Artificial Intelli gence, IJCAI'16 16, pp. 1895-1901 AAAI Press, 2016. ISBN 9781577357704.
Jungin Park, Jiyoung Lee, and Kwanghoon Sohn. Bridging vision and language spaces with assign- ment prediction. arXiv preprint arXiv:2404.09632. 2024.
François-Piersy Paty and Marco Cuturi. Subspace robust Wasserstein distances. In Proceedings of the 36th International Conference on Machine Learning, pp. 5072-5081, 2019.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning, 11(--6):355-607,2000
Duy-Tung Pham, Viet-Hoang Tran, Thieu Vo, and Tan Minh Nguyen. Mixed-curvature tree-sliced wasserstein distance. In The Fourteenth International Conference on Learning Representationss 2026. URL https: / wpwpreview.eeeeeeee/ppep. n wtr////ormmmm////// wiw.mmiiddee49wi11 https::/openzzvieeer
Michael Quellmalz, Robert Beinert, and Gabriele Steidl. Sliced optimal transport the sphere. Inverse Problems, 19510))005005, 2023.
Julien Rabin, Gabriel Peyré, Julie Delon, and Mare Bernot. Wasserstein barycenter and its applica- tion to texture mixing. In International Conference on Scale Space and Variational Methods in Computer Vision, pp. 435 446, 2011.
Michael Röder, Andreas Both, and Alexander Hinneburg. Exploring the space of topic coher- ence measures. In Proceedings of the Eighth ACM International Conference on Web Search and Data Mining, WSDM 15, pp. 399-408, New York, NY, USA, 2015. Association for Com- puting Machinery. ISBN 9781450333177. doi: 10.1 00.14452684822226853444 URL https: / o..bi.org//1..1155///
22885324.
14

---

Published as a conference paper at ICLR 2026
Mahdi Saleh, Shun Chun-Cheng Wu, Luca Cosmo, Nassir Navab, Benjamin Busam, and Federico Tombari. Bending graphs: Hierarchical shape matching using gated optimal transport. In Pro- ceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp- 11757-11767, 2022.
Tim Salimans, Han Zhang, Alec Radford, and Dimitris Metaxas, Improving gans using optimal transport. arXiv preprint
2018.
Santambrogio. Optimal Transport for Applied Mathematicians: Calculus of Variations, PDEs, and Modeling. Progress in Nonlinear Differential Equations and Their Applications. Springer International Publishing, 2015. ISBN 9783319208282. URL bttps://boone google.
com. sq www.ooooooooooece....o.
Ryoma Sato, Makoto Yamada, and Hisashi Kashima. Fast unbalanced optimal transport on tree. In Advances in neural information processing systems, 2020.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural net- works. In International conference machine learning, pp. -329-9323393 9332.. PMLR, 2021.
Charles Semple and Mike Steel. Phylogenetics. Oxford Lecture Series in Mathematics and its Applications, 2003a.
Charles Semple and Mike Steel. Phylogenetics, volume 24. Oxford University Press, 2003b.
Anthony Simeonov, Yilun Du, Andrea Tagliasacchi, Joshua B Tenenbaum, Alberto Rodriguez, Pulkit Agrawal, and Vincent Sitzmann. Neural descriptor fields: Se (3)-equivariant object rep- resentations for manipulation. 2022 International Conference on Robotics and Automation (ICRA), pp. 6394-6400. IEEE, 2022.
Shashank Singh, Ananya Uppal. Boyue Li, Chun-Liang Li. Manzil Zaheer, and Barnabás Póczos. Nonparametric density estimation under ad adversarial losses. Advances in Neural Information Pro- cessing Systems, 31, 2018.
Jascha Sohl- Dohl-Diccsttin, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing. pp. 2265-2265. PMLR, 2015.
Karen Sparck Jones. A statistical interpretation of term specificity and its application in retrieval. Journal of documentation, 28(1) 11-21, 1972.
Bharath K Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Bernhard Schölkopf, and Gert RG Lanckriet. On integral probability metrics, d-divergences and binary classification. arXiv preprint arXiv:0901.2698, 2009.
Akash Srivastava and Charles Sutton. Autoencoding variational inference for topic models, 2017.
Yuki Takezawa, Ryoma Sato, Zoritssa Kozareva, Sujith Ravi, and Makoto Yamada. Fixed sup- port tree-sliced Wasserstein barycenter. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151, pp. 1120-1137. PMLR, 2022.
Silvia Terragni, Elisabetta Fersini, Bruno Giovanni Galuzzi, Pietro Tropeano, and Antonio Cande- lieri. OCTIS: Comparing and optimizing topic models is simple! In Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguisties: System Demonstrations, pp. 263-270. Association for Computational Linguistics, April 2021.
Hoang Tran, Thieu Vo, Tho Huu, An Nguyen The, and Tan Nguyen. Monomial matrix group equivariant neural functional networks. In Advances Neural formation Processing Systems, volume 37, pp. 48625-46655. Curran Associates, Inc., 2024a. URL https::/prooceedaeee neurips wwwpeeeeeeeeeeeaeeeeeeeee--e Conferenceeeeeceeeeeeeeeeeeeeeee .df..
Hoang V. Tran, Minh-Khoi Nguyen-Nhat, Huyen Trang Pham, Thanh Chu, Tam Le, and Tan Minh Nguyen. Distance-based tree-sliced Wasserstein distance. The Thirteenth International Con- ference Learning Representations, 2025a.
15

---

Published as a conference paper at ICLR 2026
Hoang V. Tran, Thieu Vo, An Nguyen The, Tho Tran Huu, Minh-Khh Nguyen-Nhat, Thanh Tran, Duy-Tung Pham, and Tan Minh Nguyen. Equivariant neural functional networks for transformers. In The Thirteenth International Conference on Learning Representattonss ICLR, 2025b.
Hoang-Vics Saing-Viettine Tran, Thicu N Vo, Tho Tran Huu, and Tan Minh Nguyen. A clifford algebraic approach to e (n)-equivariant high-order graph neural networks. arXiv preprint arXiv:2410.04692 2024b.
Huy Tran, Yikun Bai, Abihith Kothapalli, Ashkan Shahbazi, Xinran Liu, Rocio P Diaz Martin, and Soheil Kolouri. Stereographic spherical sliced Wasserstein distances. In Forty-first International Conference on Machine Learning. 2024c.
Thanh Tran, Hoang V. Tran, Thanh Chu, Huyen Trang Pham, Laurent Ghaoui, Tam Le, and Tan Minh Nguyen. Tree-sliced wasserstein distance with nonlinear projection. In Forty-second International Conference on Machine Learning, 2025c. URL https: / //tp///vveeeeeeeeeeeeeeeeeeeeeee net/ forum?ide SWXYVIIIW..
Viet-Hoang Tran, Trang Pham, Tho Tran, Tam Le, and Tan M Nguyen. Tree-sliced Wasserstein distance on a system lines. arXiv preprint arXiv:2406. 13725, 2024d.
Viet-Hoang Tran, Trang Pham, Tho Tran, Minh Khoi Nguyen Nhat, Thanh Chu, Tam Le, and Tan M Nguyen. Tree-sliced wasserstein distance: A geometric perspective. arXiv preprint arXiv:2406. 13725, 2024e.
Viet-Hoang Tran, Thanh T Chu, Khoi NM Nguyen. Trang Pham, Tam Le, and Tan M Nguyen. Spherical tree-sliced Wasserstein distance. arXiv preprint arXiv:2503.1 11249, 2025d.
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 on Neural Information Processing Systems, 2025e.
Viet-Hoang Tran, Van Hoan Trinh, Khanh Vinh Bui, and Tan M Nguyen. On linear mode connec- tivity of mixture-of-expert architectures. arXiv preprint 1tXXi:2509.11344 2025f.
Viet-Hoang Tran, An Nguyen The, Benoît Guérand, Thicu Vo. and Tan Minh Nguyen. Quasi- equivariant metanetworks. In The Fourteenth International Conference on Learning Representa- fions, 2026a. URL https: / /openreview net/ forum?id=XMiDpp 2mWWY.
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 Learning Representations, 2026b. URL https: / //oppreevieeeee net htt/ forum? i 1d=kDgG03v055
Ananya Uppal, Shashank Singh, and Barnabás Póczos. Nonparametric density estimation & con- vergence rates for GANs under Besov IPM losses. Advances in neural information processing systems, 32, 2019.
Ananya Uppal, Shashank Singh, and Barnabás Póczos. Robust density estimation under Besov IPM losses. Advances in Neural Information Processing Systems, 33:5345-5355, 2020.
C. Villani. Optimal Transport: Old and New, volume 338. Springer Science & Business Media, 2008 Cédric Villani. Topies in Optimal Transportation. American Mathematical Society, 2003.
Thieu Vo, Hoang V. Tran, Tho Tran Huu, An Nguyen The, Thanh Tran, Minh-Khoi Nguyen-Nhat, Duy-Tung Pham, and Tan Minh Nguyen. Equivariant polynomial functional networks. In Forty- second International Conference on Machine Learning, 2025. URL https: http:s://openreiiee /owpn/eeeeeee.eeeeee net/forum wwdd===DDDDpppp
Robin Walters, Jinxi Li, and Rose Yu. Trajectory prediction using equivariant continuous convolu- tion. arXiv preprint arXiv:2010.1 1344, 2020.
Jie Wang, Rui Gao, and Yao Xie. Two-sample test with kernel projected Wasserstein distance. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151, pp. 8022-8055 8055. PMLRR 2022.
16

---

Published as a conference paper at ICLR 2026
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through align- ment and uniformity on the hypersphere. In International conference on machine learning, pp- 9929 9939. PMLR, 2020.
Wenxuan Wu, Zhongang Qi, and Li Fuxin. Pointconv: Deep convolutional networks on 3D point clouds. In Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pp. 9621-9630. 2019.
Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. arXiv preprint arXiv:2112.07804, 2021.
Makoto Yamada, Yuki Takezawa, Ryoma Sato, Han Bao, Zomitsa Kozareva, and Sujith Ravi. Ap- proximating I-Wasserstein distance with trees. Transactions on Machine Learning Research, 2022. ISSN 2835-8856.
Mikhail Yurochkin, Sebastian Claici, Edward Chien, Farzaneh Mirzazadeh, and Justin M Solomon. Hicrarchical optimal transport for document representation. Advances in neural information pro- cessing systems, 32, 2019.
17

---

Published as a conference paper at ICLR 2026
TABLE OF NOTATION
Rd 12 gd
d-dimensional Euclidean space
Euclidean norm
(d - -i-dmeensional hypersphere
unit vector
U
disjoint union
(X(X)
space of Lebesgue integrable functions on X
P(X)
space of probability measures X
M(X)
space of measures on X
p.v
measures
8())
1-dimensional Dirac delta function
U(
uniform distribution on gd-1
C(X,Y)
space of continuous maps from X to Y metric in metric space translation group of order d orthogonal group of order d tree metric
d(:,))
T(d)
O(d)
dr(:.)
((d)
Fuclidean group of order a-Waaserstein distance
W,
SW,
Sliced p- Wasserstein distance (rooted) subtree
Sp
p-Sobolev Integral Probability Metric p- p-Regularized Sobolev Integral Probability Metric tree system or spherical tree system spherical ray
Sp
T
ry
number of Monte Carlo samples
number lines system of lines or a tree system Radon Transform Systems of Lines, or Radon Transform on Spherical Trees
Ra
Ak-1
(A: -)ddimmnsion standard simplex splitting map tuning parameter space of tree systems
C.
T
distributions on (components of) space of tree systems normal (Gaussian) distribution
N
u
uniform distribution
8
Dirac delta distribution
18

---

Published as a conference paper at ICLR 2026
Appendix of "Tree-Sliced Sobolev IPM"
Table of Contents
Background on Distances on Metrie Spaces with Tree Metries
20
Background Regularized Sobolev IPM for measures Tree
20
B.1
Setting and notations
B.2 Sobolev IPM for probability measure tree
21
B.3 Regularized Sobolev IPM for probability measure on tree
22
Background on Tree-Sliced Wasserstein Distance on Euclidean Spaces
23
C.I Tree System
23
C.2 Radon Transform Tree Systems
24
C.3 Tree-Sliced Wasserstein Distance for Probability Measures on Euclidean Spaces
25
Theoretical Proofs
25
D.1 Proof for Theorem 3.2
26
D.2 Proof for Theorem 3.3
27
D.3 Proof for Theorem 3.5
28
D.4 Proof for Theorem B.3
................................
29
D.5 Proof for Theorem B.5
30
D.6 Proof for Theorem B.6
31
E Spherical Tree-Sliced Sobolev IPM
32
E.1 Background on Spherical Free-Sliced Wasserstein distance
32
E.2 Spherical Tree-Sliced Sobolev IPM
34
E.3 Properties of Spherical Tree-Sliced Sobolev IPM
34
E,4 Theoretical Proofs for Spherical Tree-Sliced Sobolev IPM
35
Experimental Details
37
F.I
Runtime and memory analysis
................................
37
F2
Gradient Flow on Euclidean space and on the sphere
................................
38
F3
Diffusion Models
40
F.4
Self-Supervised Learning
F.5
Topic Modeling....
41
F.6
Effects of the Number Trees (L) and Lines per Tree (k)
42
F.7
Effects of the Order (p)
43
F88
Analysis of Advantages of Higher Order (p 1)
44
F99
Hardware settings
45
Broader Impacts
45
19

---

Published as a conference paper at ICLR 2026
A BACKGROUND ON DISTANCES ON METRIC SPACES WITH TREE METRICS
We denote a tree as T = (V, where V and E' represent the sets of vertices and edges, respec- tively, and let r E V be the designated root. Each edge e E E is associated with a non-negative weight we that denotes its length. Following Semple & Steel (2003a), we now provide a formal definition of the corresponding tree metric.
Definition A.1 (Tree metric). Let be a and let d Q x f - [0. oo) be metric. We say that d is tree metric on f2 if there exists a weighted tree T such that f S V(T), and for any y E Ω the distance d(x, y) equal the total weight of the unique path in T connecting æ and y.
Suppose V is subset of vector space, and let (::-) denote the metric associated with T. For any two points I, E T. let [z, y] represent the unique shortest path in T connecting them. Consider the unique Borel (length) measure on T such that
w([x, v)) d v), Va,y
(13)
Additionally, given root T, we define the subtree rooted at x E T
A(x) = {v T T: : }r.
(14)
We denote by PTT) the set of all probability measures on T, that is, the collection of Borel measures on T with total mass equal to one. We now define the Wasserstein distanc on the space metric following the work of Le et al. (2019).
Theorem A.2 (Wasserstein on Tree Metric Spaces). Let u, P(T). Then the 1-Wasserstein distance with respect to the tree metric dr can be expressed as
W1,t (p, v)
((l(A(()))
w(dr),
(15)
where A(x) denotes the subtree of rooted at x. and w is the associated length measure on T.
While the I-Wasserstein distance (W,) tree has convenient closed-form solution, this is gener- ally not true for higher orders (p > 1). This computational challenge motivated the development of alternatives like Sobolev Transport (ST), which provides scalable and valid metric for comparing probability measures on tree and graph structures (Le et al., 2022).
The key idea behind ST is to modify the constraints on the "critic" function used differentiate between two measures. Instead of the standard Lipschitz condition, it constrains the critic function within graph-based Soholev space, primarily by limiting the T.Pnormm of the function's gradi- ent (Le et al., 2022). This approach has proven versatile, with extensions for measures of different total masses (Le et al., 2023) and for more general geometric structures beyond the standard L framework (Le et al., 2024).
closely related concept is the Sobolev Integral Probability Metric (Sobolev IPM), which is a type of IPM where the critic function is constrained to a unit ball defined by the full Sobolev norm measure that considers both the function's values and its gradient (Adams & Fournier, 2003). The crucial innovation is a regularized variant of this metric. By relaxing the constraint to focus only the gradient of the critic function, the regularized Sobolev IPM successfully admits a closed-form solution for any order p > 1, making it a powerful and computationally efficient tool for comparing measures on trees (Le et al., 2025).
B BACKGROUND ON REGULARIZED SOBOLEV IPM FOR MEASURES ON TREE
In this section, we introduce the framework of regularized Sobolev integral probability metrics (IPMs) for probability measures supported on tree structures. Specifically, we begin by review- ing the tree setting for probability measures, including the relevant notational conventions and the functional spaces. We then formulate the Sobolev IPM problem for measures supported trees. Fi- nally, we introduce the regularized variant of the Sobolev IPM, which admits a closed-form solution. Proofs in Section closely follow that of Le et al. (2025).
20

---

Published as a conference paper at ICLR 2026
B.1 SETTING AND NOTATIONS
Let T = (V, E) be a finite rooted tree, where V denotes the set of vertices and E the set of edges. For each edge e E, we associate positive weight w. The tree T is naturally equipped with a metric dy (z v)
which represents the length of e. x V - I such that for x, E V. Eee(nag) We is the length of the unique path connecting and y. Furthermore, given two vertices x, 2 E V, we let [x, 2] denote the path connecting x and 2. Moreover, for each edge E E, we denote by ve the endpoint of that is farther from the root of TT, and by % the subtree of T rooted at v.c.
Measures and functions. Let P(T) denote the set of all nonnegative Borel measures and let (T x T) denote the corresponding set of measures on the product space T x T with finite mass. A function f: T - R is said to be continuous if it is continuous with respect to the topology on T induced by the Euclidean distance. We write C(T) for the space of all continuous functions on T. and analogously C(T 7) for continuous functions on T T. Given nonnegative Borel measure w on T and an exponent 1 S p < co, we define the space L" (w) as
graph T,
IP(T,W) w) = { {T:T-R
I (1(6)) w(dx) < oo).
(16)
This is normed space equipped with the norm || (T11(T)) (
In addition, let Φ - R.+ be strietly positive weight function, i.c., w(x) 0 for every a E G. The associated weighted TPP space, denoted T.e (G,w), is given by
L (T,w) { ( - / ((o)))
(17)
B.2 SOBOLEV IPM FOR PROBABILITY MEASURE ON TREE
Following the definition of graph-based Sobolev spaces Le et al. (2022), we define the tree-based Sobolev space as follows.
Definition B.1 (Tree-based Sobolev). Let be nonnegative Borel measure on the tree T, and let 1 S p < co. A continuous function ʃ T - IR is said to belong to the Sobolev space W1.p (T, w)
there exists function h E L"(T, such that
[ [sooee)
f(x) f(20)
h(y)
Va ET,
(18)
where 20 is fixed reference vertex in T. The function h is uniquely determined in LP(T, w) and is referred to as the tree derivative of with respect to the measure w. We denote the tree derivative of WLP(( (w) by f' LT((T,,) ww)
The Sobolev space W WL.P(T, w) is endowed with the norm
I = (I1S + I1 1/p
(19)
which is referred to as the Sobolev norm (Adams & Fournier, 2003) turning W w) into a
normed space. We further denote the subspace W w) {f E WT.((T,w) f(20) = 0}. T. where 2n is the root of T. The unit ball in this space then denoted by
B(p) = {f E W.
w) : Ilfl $1}
(20)
Definition B.2 (Tree-based Sobolev IPM). Let w be nonnegative Borel measure on T. and let
1 S p S oo with conjugate exponent p' defined by 1 + + = (with the convention p' = oo when
P p 1). For two probability measures μ., P(T), the Sobolev IPM is defined as
/ (())))((a V) = sup SEB(Y)
f f(x) v(dy)
(21)
where B(y') denotes the unit ball in W!!-r' (T,
21

---

Published as a conference paper at ICLR 2026
Overall, the Sobolev IPM for probability measures graph be viewed particular case of the IPM, where the witness critic functions are restricted to the graph- based Sobolev space and further constrained to lie within its unit ball. Furthermore, Notice that the quantity inside the absolute signs is unchanged if f is replaced by f - [(zo). Thus, we can assume without loss of generality (T,w). Next, we introduce a weight function. Let w that (zo) O. This is the motivation for W. (T,((TW) be measure w on set. We have
( = : 1=)==1+(AAxx)) E
(22)
An example of w is when it is chosen as the length measure (Le et al., 2022), in which case (A(z)) corresponds to the total length of the subtree A(x). We now present key theorem establishing that, for any critic function f E Wo (T, w), the Sobolev norm is equivalent to the weighted normorm of its derivative f".
Theorem B.3. Let w be nonnegative Borel measure on T and let 1 p < co. Define the constants
(1= (minn(1,w(T)r 01 == (1+ (T)
Q2 =
(23)
Then, for every f E (T.""(T.s)) the following norm equivalence holds:
a ll S IIflww. 02 ll
(24)
Proof of Theorem B.3 is defered Section D.4.
B.3 REGULARIZED SOBOLEV IPM FOR PROBABILITY MEASURE ON TREE
Having established the equivalence relation in Theorem B.3, we now introduce the regularized Sobolev IPM. Specifically, rather than constraining the critic ʃ to lie in the unit ball B(y') of the Sobolev space, we instead restrict ʃ to the unit ball Bip', vi), defined with respect to the weighted I. -norm of its derivative f' under the weight function w. Hereafter, define B(p', ()
B(p () == { Wo (w,w) 11f' $1]
(25)
We now formally define the regularized Sobolev IPM between two probability distributions on tree T
Definition B.4 (Tree-base Regularized Sobolev IPM). Let w be nonnegative Borel measure T and let 1 S S co. For any probability measures pl, E P(T), the regularized Sobolev IPM is defined by
(.aaataa (())(dd)))
up(u,))==
sup JEB(y' .i)
(26)
w where B(p', w) denotes the unit ball in the weighted Sobolev space induced by the norm (r'
Next, we show that the tree-based Sobolev IPM has a closed- solution that is as follow
Theorem B.5 (Closed-form Expression). Let w be nonnegative Borel measure T, and let p < sc. Then, for any probability measures μ1, u E P(T), the regularized Soboley IPM admits the closed-form expression
Sp = ((xx)) 1-p ((()))
(27)
where (xx) denotes the subtree rooted at x.
Proof of Theorem B.5 is defered Section D.5. Additionally, when the input probability measures are supported on nodes V of T and we choose the length measure on tree for the nonnegative Borel measure w,
derive explicit formula for Equation (26) as follow
Theorem B.6 (Explicit formula for Discrete Case). Let w denote the length measure on T, and let S DC. Suppose that p. P(T) are supported on the vertex set V of the tree T. Then the regularized Sobolev IPM admits the closed-form expression
1/p
Σ Be (p(Te) -((e)p
(28)
NEEB
22

---

Published as a conference paper at ICLR 2026
where for each edge E, the scalar weight Be is given by
We
log
1+ 1 w(7
2,
Be
(29)
(1+w(1e)) we p (1 (
otherwise,
2-P
with we denoting the length of edge e and Te the subtree rooted at the endpoint which is the endpoint of edge e farther from the root.
Proof of Theorem B.6 is defered to Section D.6.
Implementation of Explicit Form for Discrete Case. The regularized Sobolev IPM S, depends only the graph structure (V, E) and edge weights {weJee v. and therefore be applied beyond physical graphs. For efficient computation, the sets and coefficients 3e (cf. Equations (10)-(1)) can be precomputed once from the root 2g to all vertices.
C BACKGROUND ON TREE-SLICED WASSERSTEIN DISTANCE ON EUCLIDEAN SPACES
This section revisits the fundamental components of the Tree-Sliced Wasserstein (TSW) distance, formulated tree systems embedded in Euclidean spaces. For completeness, summarize key definitions and core mathematical formulations. The reader is referred to (Tran et al., 2024d; 2025a) for detailed proof' and detailed explanation.
C.1 TREE SYSTEM
Figure 2: The construction of the tree system is illustrated in the two-dimensional plane R2, though the approach naturally extends to higher dimensions. The procedure begins with set of infinite lines placed without any predefined arrangement. All pairwise intersections of these lines determined, though some may lie outside the visible region of the figure due to their unbounded extent. Among these intersections, subset is marked in red indicate those that will be discarded. The remaining intersections in green are preserved in order to impose a tree structure on the system-ensuring that any two points along the lines are linked by unique path passing only through the retained intersections. These preserved points serve as the fundamental nodes defining the tree topology. Once the red intersections are removed, the resulting network forms the desired tree system.
Components of Tree Systems. line in Fd is an element (x, 0) E 12d $d-1 where the source and 0 is the direction. It is parameterized by x + t0 for E R.
Given an integer A: 1, system of k lines in Rd refers a collection of k such lines. The notation (Rd x gd is abbreviated as Lg representing the space of systems of i lines in Rd. An element in this space, commonly denoted by L. corresponds to a specific system of lines, written as f= }/_ where each I; = (x,, 0.) R x gd-1 and indexes the lines.
line system L is connected the union of points lying on its individual lines a connected subset of Rdd A tree structure can be imposed by removing selected intersection points so that any two points in the resulting configuration are joined by a unique simple path. The resulting object a tree system, denoted T {4,)==1- We use the term free system emphasize this unique-path property, in direct analogy with trees in graph theory. Using remaining intersections, we build a topological tree system by coherently gluing segments of R via disjoint union and quotient topology (Hatcher, 2005), resulting in a space endowed with a valid tree metric. An illustration of the construction appears in Figure 2.
23

---

Published as a conference paper at ICLR 2026
Sampling Tree Systems. Tree system spaces admit diverse structures, but Tran et al. (2024d) high- light chain-structured variants. A generative model for such systems is as follows: sample an initial point 21 M4 and direction 01 v1; then for each i >> 1, sample offset t μ4 and direction 0 ~ Vi, and set 3i = - + 1,0;-1- Rach μ4 and v; is an independent distribution. In practice, we take 1a1 = (1-11,11))) 11 pr ((--1,1) 1J) for i 1, and v4 (S for all i.
In Tran et al. (2025a), a concurrent-lic tree structure is introduced, where all lines (4)]-z share the same source point æ. The corresponding generative model is simpler: first sample the common root ~ then independently sample 0; ~ v for each i = 1..... k. Here U(I-1, 1J9) specifies the distribution of the root, and // ((( the common distribution over directions.
A visualization of the two tree structures is provided in Figure 3.
Remark C.1. Recent advances in Tree-Sliced Wasserstein distance (Tran al., 2025a:d) employ the concurrent-line tree structure. For TS- TS--oobolev we likewise adopt this structure for its simplicity.
Remark C.2. The chain-structure and concurrent sampling schemes each induce probability measure a over the space T of tree systems.
Figure 3: Visualizations of two popular tree structures: a chain structure (left) and a concurrent-
structure (right).
C.2 RADON TRANSFORM ON TREE SYSTEMS
Let L (RR) be the space of sebesgue-interriie Fix a tree system T { E Tr. with lines 4 = (æ,, 0z), and let T denote the union all points
functions Rd, equipped with the norm || 1h--
on these lines. A function T - I is integrable over T
EV ((tz,14) dt < co,
Ifll
(30)
and the collection of such functions is denoted L L'(T). the space of Lebesgue integrable functions over the tree system T. The standard 1)- simplex
k O,
A+-1 (a= ax) E R
a 2
Σ a. = 1
(31)
Write C(R. Tg. A for the set of continuous splitting maps a R
Ag-1-
Given L' (Rd) and a E C(R2 Ty. Ak-1), we define the projection operator that maps 1 (Rd) a function on T:
Rq R9,(244 S f(y)
a(y, T)x Ott - 0.)) dy,
(32)
24

---

Published as a conference paper at ICLR 2026
where (z,, 8,) specifies line and ó is the Dirac delta. According Tran et al. (2025a), R7 (T) and /|TT & || hence Rg. (. (T) is well-defined operator.
Aggregating over all tree systems T E TV we define the Radon transform on tree systems by
R (RR)
II (T).:
(33)
TET
f (R TET
(R
(34)
If is invariant under the Euclidean group E(d), then R is injective.
C.3 TREE-SLICED WASSERSTEIN DISTANCE FOR PROBABILITY MEASURES ON EUCLIDEAN SPACES
Consider probability measures J4, E P(RR). Given tree-structured collection of lines T E T and an E(d)-invariant splitting map a C(TR x Lg Ak-1), )) the operator R.j. transports jr and onto corresponding measures R and R4 supported on T. Since each tree system T is endowed with tree metric dT, one can evaluate the 1-Wasserstein distance Wa..1 between these transformed measures. This motivates the definition of the Distance-based Tree-Sliced Wasserstein (Db-TSW) distance Tran et al. (2025a):
DD-TSW v) = L Wi (R9 Rgv) do(
(35)
where a denotes probability distribution over the space of tree systems T. The value of Db-TSW depends on the choice of tree space T, the sampling process that induces a, and the B(d)-invariant map a, although this dependence is suppressed in the notation for simplicity. The resulting Db-TSW provides an E(d) -ivvariatt metric on ((..).
Remark C.3. When the tree systems consist of single line, the Db-TSW distance reduces to the standard Sliced Wasserstein distance.
Splitting Maps. Since R" is injective whenever the splitting map is E(d) ((dd)-nnvriian, we seek constructions that satisfy this property. For E IRd and a tree system T (4,)2 Ty. define the Euclidean distance from to the line 1, E T TTby by
d(r, T) inf || 1am-y/i =112, yEl.
(36)
which invariant under the Fuclidean group F(d).. Any splitting rule that depends only on the collection (d(x, T therefore inherits F((d)invariancee. A practical and widely used choice is the softmax:
alx,, T) = softmax((-$ d(x, )
(37)
where the parameter & > controls the sharpness of the distribution. This map assigns weights to the lines in T according their proximity xr while preserving Euclidean symmetries. Empirically, softmax-based splitting maps have been found to perform well in applications (Tran et al., 2025a).
Remark C.4. Equivariant neural networks (Cohen Welling, 2016) explicitly encode symmetry priors of a task directly into the architecture, thereby promoting better out-of-sample generalization and higher sample efficiency. They have achieved substantial empirical success across diverse areas, including trajectory forecasting (Walters et al., 2020), robotics (Simeonov et al.. 2022), graph repre- sentation learning (Satorras et al., 2021; Tran et al., 2024b). Optimal Transport- driven frameworks (Tran al., 2024e; 2025e), equivariant metanetwork design (Tran al., 2026a: Vo et al., 2025; Tran et al., 2024a; 2025b), and analyses conducted in weight space (Tran et al., 2025f).
D THEORETICAL PROOFS
In this section, we prove results for TS- TS-Sobolev stated in the main text and establish the properties of Sobolev IPM introduced in Appendix B.
25

---

Published as a conference paper at ICLR 2026
D.1 PROOF FOR THEOREM 3.2
Proof. We show that TS-Soboley, is an E(d))invariant metric on the space of probability measures P(RR). The definition is given by
TS TS-Sobolly,,(,, = (1, Sp ur) d "
(38)
where S, is the regularized Sobolev IPM, which a metric on the space of measures on tree T. The proof relies on the injectivity of the Radon Transform Ro, which holds because the splitting map a is chosen to be F(dd)invariiat (see Tran et al. (2025a)).
We now verify the three metric axioms.
Positive definiteness. is clear that TS-Sobolev, μ)) = O and TS-Sobolev, 2 TS- (i-Soollee,p(1.11 u) = 0, this implies Sp(T, "T) dT( = 0. Since the integrand is non- negative, this means Sp(uT, UT) for almost MT MT for almost all T. By the injectivity of Ra conclude that the densities equal, fu = fu. and thus v.
As Sp metric, follows that
Symmetry. The symmetry of S, on each tree implies Sp(NT, (p(UT,HT)P. Therefore,
(L, ST(TT,
(S-Sobolly,((u.))
(L. Sp ((( (TS
(39)
Triangle inequality. For any ,1, U2, 14 E P(RR), we use the triangle inequality of S, on each tree, which states (,,.T & H,,, + Ha,T) We then apply Minkowski's integral inequality:
Sobolev,(uu. H3) ( Sp (A.T) do(T)
(Sp(NT (2,T) + (((22TT.HTT)) do(T)² :
do(7) $
( Sp(un. Pod))
+(, (o do(T))
(1--ooolle,,(1,1,1 12) + TS-Sobolle,((22.11))
(40)
Thus, Toboley, is a metric on P(Rd).
invarrineee We aim to show that for any E E(d), TS-Sobolev, TS-Sobolev, gtv). Let T = (li = (r,, 0 be a tree system. Under the action of = (Q, v), we have gT = {g4 = (Qx; uu Q0) Since dde((Q) = and is E(d)-
invariant, we compute for a line 4 T:
[ ((.. (y- Igli )gt))) dy
[ Sulo al 6(1 - -y
(y- - Fgl,, 0gt,) dy St.. fulg fu (alg~'y) ag)a(g...00 9T): 8(tz (gy Igl. )z..)) d(gy)
[ (p(v) ol, T): 0tz (ov Fghh Ogt,) )) dy
26

---

Published as a conference paper at ICLR 2026
J ful)) (Qy+ (Qx; + +), Q0.)) dy
[ (ule)a(u. T): o(z (Q(v - mz)) Q0.)) dy
[ f,l))a(y. T) (ts (v 1-=1, zz. 0.)) dy
(2,4(244)
(41)
This implies that the action of is an isometry, so Sp(RT, VT) = (0 Using this, we compute:
TS-Soboley, gtv)p (allatu), (T
f Sp((otu)gT (oth)) do(gT)
J (, do TS-Sobbole,,uu.1)
(42)
Taking the p-th root of both sides, we conclude that TS-Sobolev, is E(d)-invariant.
Remark D.1. For clarity, we omit the almost-sure conditions in the proof. Verifying these condi- tions is straightforward, and their inclusion would make the core argument harder follow.
D.2 PROOF FOR THEOREM 3.3
Proof. We prove the theorem in two parts. First, establish the equality for the case 1. Second, we prove the general inequality for any E [1, z0).
Part 1. We first proof equality for 1. By definition, the Tree-Sliced Sobolev IPM is given by:
(s-Sobolev,((1.)) Sp do(e))))do())))) do(())
(43)
For the case p 1, this definition simplifies to the expectation of the base metric:
(s-Sobolev.(,,1) = f
(44)
The Tree-Sliced Wasserstein distance is defined TSW v) - WI(T, (T). To prove the equality, it sufficient show the integrands are equal, i.c., Si(HT, UT) = (I(T,,T)) We analyze the discrete form of the Regularized Sobolev IPM from Equation (5) for p 1:
(1 PT)
Σ
Be (u(TT)) - v(Te)l.
(45)
BEE
From Equation (6), the coefficient Be simplifies to Be we for p 1. Substituting this result gives:
SI(UT, VT) = Ewe Iu(T)) v(Te))
(46)
This expression is the known closed-form solution for the 1- Wasserstein distance on a tree, W1 MT) (Le al., 2019). Since the integrands are equal, their expectations equal, which proves that TS- S3-Sobolev, = TSW (W) v).
We now prove the general inequality TS-Soboley,,Hee w)2 < TSW( v). It is sufficient to show that the integrand the first expression is bounded by the integrand of the second on any given tree T. That is, we will prove (, " < Wi Wi(UT,YT)
Part 2.
We first establish that the Sobolev coefficient Be w. for all p by case analysis.
27

---

Published as a conference paper at ICLR 2026
Case 1: p = The coefficient is 3e = log [1 (4w(re))
Using the inequality log(1+ + +2) I x for
O, we have:
Be = log 1 1 w(T))
We
We
(47)
+ ((%e)
Since w(Y))= 2 O, the denominator at least 1, which implies I+w[T)] I we- Thus, 3g I W/-
we
Case 2: p 1 and p o We apply the Mean Value Theorem. Let the function be f(x) x² and consider the interval [a, bj where a = 1 1 w(%) and = 1 + w((e) + We. The theorem states there a value e E (a. b) such that f(b) f(a) (c) a). The derivative fc( (2 p -P Substituting this into the definition of pe:
Be + (1+e(77) + wa)2- (1 + w(To))
--p
f(b) f(a) (b a) p)ppcc 2-p
We cPPP.
(48)
2-p
Since = + w(Tr) 2 1, the intermediate value e must greater than 1. For any p 2 1, the exponent 1 is non-positive (5 0), which ensures cl < 1 and therefore Be S we.
Now, using the universal bound 3. I We and the fact that (u():) v(7e)l E [0, 1], which implies I... I.. I.... ... for 1, we can bound the p-th power of the Sobolev IPM:
S uT)
Σ eEE
Be (m(T)) u(T)))P
Σ
We (u(T)) v(Te)lp
(since Be S we)
E
Σ eEE
We (u(T)) v(7)))
(since |...| 1' s I....
Wi(uT,uT)
(by Part I of this proof)
(49)
This establishes the key inequality on a single tree: (p(HT,UT) S WI(HT, Integrating this inequality over all trees T T directly yields the theorem, as the left-hand side becomes TS-Soboley, v and the right-hand side becomes TSW( This completes the proof.
D.3 PROOF FOR THEOREM 3.5
Proof. We analyze the convergence of the Monte Carlo estimator for the TS-Sobolev, which is defined as:
(1 7
L S. Sp(HT, vT))")
-
TS Sobolee (u,v)) =
(50)
Let us define the random variable X; = SpT, (T,) where each X; is an independent sample drawn by sampling mean X Σt= X,.
T. a. The estimator can then be written as a function of the sample
The expected value of X; is the quantity we are trying to estimate (raised to the p-th power):
px = E[X.] = { S, do TS-Soboley, v)P.
(51)
Let the variance of X, be finite, denoted by ox = V[X,] By the Central Limit Theorem, the sample mean X is asymptotically normal, and its variance is V X /L.
Our estimator is function of this sample mean, specifically (XX), where g(y) // To find the variance our estimator, we apply the Delta Method. The variance of g(X) can be approximated by:
v[g(X)) (ux))) v(x
(52)
28

---

Published as a conference paper at ICLR 2026
where s' (uxx is the derivative of g evaluated the mean Jax. The derivative is 5'(()
Substituting this into the variance approximation, we get:
V/ (1s- w)] ~
(- L
+
xx
(53)
The Root Mean Squared Error (RMSE) the estimator is the square root of the variance.
RMSE vv (TS-Soboleyy((..)) VT.
1
ox.
(54)
p
Since the terms jax and Øx finite constants that do not depend on L, the Monte Carlo approxi- mation error decays at the standard rate of OLL-1/2).
D.4 PROOF FOR THEOREM B.3
To ensure a rigorous and self-contained presentation, we now derive the result in full, adopting the framework proposed by Le et al. (2025).
Proof. Let f (T,"), We first derive an upper bound for II in terms of [| Since
f(zn) = 0, it follows that
Il ( [[(())) w((x)
= =fr Mm ((()) " w(dr)
l 1r (emne)()) ('v) u(d))) (d))) w(dr).
Applying Jensen's inequality, we obtain
I (w(P)= J J 1 [o0 (5) w(())) w(dr).
By Fubini's theorem, we may interchange the order of integration, which yields
I (W())-1 [ |(() (fr 1/00.)1000 (0) w(dr) w(ly)
= / Pf ((T()))
(55)
where recall that T(y) = for E T v E [zo 2|]. Using the estimate from Equation (55), obtain
(1FIIEP + || =
(WIT))-) ( [1+ta1 + (+()) 1
((AA)) (())) ((
1
az 1l
(56)
where r6(x) = 1 (>(x))) and 02 1/p
29

---

Published as a conference paper at ICLR 2026
Next, we derive corresponding lower bound for Ilflw Since || 2 0, it follows that
I = (1f1l2> + 11(())))
2 1l
- (f (((r)))(((a)))
I (T)P) (T)P]] ( (()))) "/
1
2 (min(l,
f ([ w(do)
(+)(T))
= an |1
(57)
where
(min( 1/p
Q1 ( wT)p
(58)
Combining Equation (56) and Equation (57), we conclude that @1 llf'lleg & ||llwir a2 ll
(59)
which completes the proof.
D.5 PROOF FOR THEOREM B.5
To ensure rigorous and self-contained presentation, now derive the result in full, adopting the framework proposed by Le et al. (2025).
Proof. Let f E W/ (T,w). By Definition B.1, have
f(z) = f(20)( + f [20.>]
w(d)), VIET.
(60)
Using Equation (60) together with the indicator function of the path [zo, æ/.. and noting that ((T)
1, we obtain
1 ((())((d)) 1 f(20) p(dx) +
Sr/ (cn.2)) (')) u(dx)
=((2)) fr// (/000) (v) f'()) w(dy) p(dx).
Applying Fubini's theorem to interchange the order of integration yields
{ ((o)))l(d)) f(zo) [ (// 1 (=m (v) u(dx)) f'(() (v) w(dy).
By the definition of [(y), this becomes
f (()) f(2o) P(T))())))))))))
(61)
An analogous computation gives
/ [())v(dz) -((0) + (()
(62)
30

---

Published as a conference paper at ICLR 2026
Hence, the regularized Sobolev IPM Equation (26) can be written as
FEB(p',0) .0) J,, r(w) -
bp(u,D= = sup
((A(z))) w(dx)
(63)
where
B(p',0) (f Wo (T,w): 1f :T $ 17.
(64)
Observe that
: E ()) (g L' (T,w) ||91| og' 11..
(65)
Indeed, the inclusion "S" is immediate, while the reverse direction follows by constructing f(z) := JFo... (3) (dy) for any E L (T,w).
Now define
f(x) p(A(z))- -(A(x))
TET.
(66)
(x))
Substituting into Equation (63), we obtain
V ((xx) f(x) o(a))wwd)))
Sp(1) sup gELN (T
(67)
= ( (()) ( 1/0
(68)
{f (U))) w((s))))()())) 1/p
where Equation (68) follows from the dual norm characterization of weighted Lt spaces.
Therefore,
Sp [((A)))))) w(dx),
(69)
which proves the claim.
D.6 PROOF FOR THEOREM B.6
To ensure a rigorous and self-contained presentation, we now derive the result in full, adopting the framework proposed by Le et al. (2025).
Proof. We work with the length measure w on the tree T, so that w(fx)) for all æ e T. From Theorem B.5, it follows that
Sp(m,
Σ
- ((A(()))) w(z)'
(70)
(E(E)EE (u.v)
Since μ, are supported on vertices V. for any E (u, v) we have
(((x)) v(A(z)) ((A(x) \ (u,w) v(A(x) (u,u))
(71)
Substituting this into Equation (70) yields
Sp(,,DP== Σ
((r)+~P| (n(A(z)
\ (u,w)) ((((a))))(((()) v(A(x) \ (u, v))) w(dx). w((d))
(72)
(=(W,))E (0,0)
31

---

Published as a conference paper at ICLR 2026
For any edge (u, o), it follows that A(x) \ (2, v) % for all (u, v). Hence, Equation (72) simplifies to
E (m(Te) - u(>e))2 [ (n(x)'-w((())
Sp(,,)==
(73)
e=(uN)EB
We now compute the integral term. Recall that w(x) (((x))) Without loss of generality, assume d u) I (20, v), i.e., is farther from the root For a E (u, v), write x + t(u - u) with (0, 1). Then
w({()) w(%)) + wet,
(74)
and therefore
J(..)) ((x)) ( w(x)) 1/' 11 w((e) wet|" 1-p we dt.
(75)
This integral evaluates explicitly as
Wr
[ ( -
log 1+ 1 + w(T))
p=
(76)
(1+)()))) + We
P- (1 (())
p * 2.
2 --p
Thus, f(u,0) (P = Be (see Equation (29)). Substituting into Equation (73), we obtain
1/p
Sp(M,D= = Σ
Be -
(77)
NEER
This proves the result.
E SPHERICAL TREE-SLICED SOBOLEV IPM
E. BACKGROUND ON SPHERICAL TREE-SLICHE WASSERSTEIN DISTANCE
In this section, review the concepts of Spherical Tree Systems, the Spherical Radon Transform, and the Spherical Tree-Sliced Wasserstein distance, as proposed by Tran et al. (2025d). These are the spherical analogs to the Euclidean framework of Tree Systems and their corresponding Radon Transforms (Tran et al., 2024d; 2025a). We will follow the construction of the Euclidean back- ground, explaining the spherical components in the same order.
Hypersphere. The underlying space for these spherical constructions is the d-dimensional unit hypersphere, denoted sd. This the set of all points in (d + 1iddimensional Euclidean space Rd+1 that are at distance of 1 from the origin. Formally:
gd {x E RR++1
||<|||= =}
(78)
This space is metric space equipped with the geodesic distance, which is the shortest distance between two points along the surface of the sphere. For any two points a, E Sd, this distance is the angle between them, calculated
(g, b) arccos ((a,b))a+))
(79)
A key tool for relating the curved geometry of the sphere flat Fuclidean geometry is the stereo- graphic projection Pa.- For point & E Sd (the "pole" of the projection). the map 9a, projects point y from the sphere onto the hyperplane H that is tangent the sphere the antipode -r.
Px Sd
{x}
H2
(I, y)
1
(80)
(x.
(ar, y)
y.
By convention, the map is extended to the entire sphere by defining ia (x(x) = oc, which completes the mapping p gd H2 U (oc).
32

---

Published as a conference paper at ICLR 2026
Components of Spherical Tree Systems. The fundamental building block for spherical tree is the spherical ray, the analog of a straight line in Euclidean space. is formally constructed using the inverse stereographic projection. The intuitive idea is to first draw straight ray the flat hyperplane H and then use the inverse map oz to trace this path back onto the curved surface of the sphere.
Mathematically, a spherical ray with root x and direction y, denoted r... is defined as:
P (tt y : 12 (
(81)
Here, the set {t- y t 0} represents straight ray on the hyperplane Hz, and of maps this line back to the sphere. Each resulting spherical ray is isomorphic to the interval [0, z via the geodesic distance from its root, dga (xx, +), which allows any point the ray to be uniquely parameterized by its distance from x.
A spherical tree system T is formed by gluing set of k spherical rays {r ) at their common root I. This construction ensures that any two points in the resulting configuration are joined by unique simple path, endowing the system with a valid tree metric, dr. The space of all such spherical trees with & rays (or edges) in gd is denoted by TV This space is equipped with a probability distribution that governs the tree sampling process.
Sampling Spherical Tree Systems. To sample spherical tree system, employ the generative process described in Tran et al. (2025d). The process begins by sampling common root æ from the uniform distribution on the d-sphere, μ = (sd). To generate the k orthogonal direction vectors, second step is performed: for each direction, an initial vector is also sampled from ((*).. This vec- tor is then projected onto the hyperplane H₂ (which is orthogonal to the root x:) and re-normalized.
Spherical Radon Transform on Spherical Tree Systems Let (Sd) be the space of Lebesgue- integrable functions on the hypersphere $d. For given spherical tree system T E Ty. we de- fine L (T) as the space of integrable functions T R such that their norm, || Σ 1o (t, dt, is finite.
A splitting map is continuous function E C(S x TV. 1 that assigns weight distribution to the rays of tree for any given point on the sphere. Given such a map, the Spherical Radon Transform on Spherical Trees projects a function f E I, (S4) onto function f in T. For point (t, r on the i-th ray of the tree, the transform is defined as:
f4,(t, - f(0) ay, Th 8(t arccos (z, v)) dy,
(82)
where d is the Dirac delta. This operator is well-defined from (S())) to (T()). Aggregating over all trees, the transform R maps function on the sphere to collection of functions on all possible trees. the splitting map a is invariant under the orthogonal group O(d + 1), this transform is injective (Tran et al., 2025d).
Spherical Tree-Sliced Wasserstein Distance For probability measures V E P(Sd), ), the oper- ator R2 transports them to corresponding measures on the tree, R9 and Rjv. Since T has tree metric dy, we can compute the 1- I-Wasserstein distance between these projected measures. The Spherical Tree-Sticed Wasserstein (STSW) distance is then defined as the expectation over all trees:
/ w (I RFv) dT
STSW
(83)
When the splitting map is chosen to be O(d + 1) -iivvariant, the STSW distance is O(d 1)- invariant metric on P(Sd).
Splitting Maps. The invariance of the metric relies on an O(d + invvariia splitting map. The group O(d + 1) consists of transformations on dd+1 that preserve the Euclidean norm and thus leave the sphere gd invariant, A splitting map a is O(d + i)-invariant a(gy.g7) = a(y. T) for all g E O(d 1).
33

---

Published as a conference paper at ICLR 2026
A practical way to construct such map is base it on an invariant quantity. For point y S and a spherical tree T, one can define an invariant "distance" " Bly, 7); from y to each ray of the tree. A continuous and O(d + 1-iinvariant choice for this map, presented in Tran et al. (2025d),
is given by 3: $ x T - IR::
0,
if y =z,
B To De
arccos
(p-arc)
VI (z, y)?
y + tz.
(84)
(x-(x,))
A valid splitting map then be constructed by applying the softmax function to these values:
aly, T); softmax ((= B(v, ())))))))i)
(85)
This map assigns higher weights to the rays that are "closer" to the point y, while preserving the necessary rotational symmetries.
E.2 SPHERICAL TREE-SLICED SOBOLEV IPM
The Tree-Sliced Sobolev IPM framework can be extended from Euclidean spaces to measures de- fined on the d-dimensional hypersphere, J.v P(S"). This is achieved by replacing the Eu- clidean components with their spherical analogs. Instead of projecting onto tree systems formed from straight lines, we project onto spherical tree systems built from spherical rays.
The projection is performed by the Spherical Radon Transform (R"), which maps the spherical measures J4, v E PP(S) to corresponding measures HTT == R. (p) and VT (2) on given spherical tree T. We then compute the regularized Sobolev IPM, Sp LT), between these pro- jected measures. The final distance is the expected value of this quantity, taken over the space of random spherical trees T with respect to a probability distribution a.
Definition E.1 (Spherical Tree-Sliced Sobolev IPM). The Spherical Tree-Sliced Sobolev IPM of order p [1, denoted as STS poboley... between μ., E P(B") is defined by
ST)-Soboley,,(..)) = SoSue.. uT) do(t))
(86)
E.3 PROPERTIES OF SPHERICAL TREE-SLICED SOBOLEV IPM
Metricity of STS-Sobolevp. The metric properties of STS-Sobolev are guaranteed by the invariance of its components under the relevant symmetry group for the sphere: the orthogonal group O(d+1). This is the group of distance-preserving linear transformations (rotations and reflections) in Rd+11 which leave the sphere 3d invariant. An O + -iiinaariant splitting map ensures that the Spherical Radon Transform is injective, which is crucial for the metric properties. Just as in the Euclidean case, this invariance guarantees that Coo--Sooleev is not only invariant but also a valid metric.
Theorem E.2. The STS- STS-Sobolev is O(d + 1)- invvarriant metric P(S)).
The proof is analogous to that of Theorem 3.2.
Connections to STSW. The STS-Sobolev IPM is natural generalization of the Spherical Tree- Sliced Wasserstein (STSW) distance (Tran et al., 2025d). It recovers STSW exactly for the case p 1 and is bounded by it for all other orders.
Theorem E.3. For any μ., E P(()) and p 2 1: (1 < STSW(H, v), with equality if = 1, iee, STS-Sobolev (u, = (p,
The proof is analogous to its Euclidean counterpart in Theorem 3.3.
Computation of Spherical Tree-Sliced Sobolev IPM. The integral in Equation (86) is intractable and is approximated using Monte Carlo estimate by sampling L spherical trees:
STS- STS-Soboley,
Σ
Spm, vTT)P)
(87)
34

---

Published as a conference paper at ICLR 2026
The computational complexity of STS-Soboles matches its first-order counterpart, STSW, (Ln log n + Lkdn). The additional step of computing the coefficients Se per Equation (6) in- troduces negligible overhead of only O (LLn). Crucially, key advantages are preserved: this com- plexity holds for any order E [1, oo), and the empirical runtime remains nearly identical to that of STSW.
Theorem E.4. The approximation error of STS-Sobolev decreases a rate of OL
The proof is analogous to that of Theorem 3.5.
E,4 THEORETICAL PROOFS FOR SPIERICAL TREB-SLICED SOBOLEV IPM In this section, we provide the proofs for the results stated in Appendix E.3.
Proof for Theorem E.2
Proof. We show that STS-Soboley, is an O(d ii-inaariant metric on the space of probability measures P(Sd). The definition is given by
ST)-Soboley,(1,,)) = SS uT do(T))
(88)
where ji and VT are the projections of μ and via the Spherical Radon Transform, and S, is metric on the space of measures on a spherical tree T. The proof relies on the injectivity of the Spherical Radon Transform R which holds because the splitting map is chosen be O(d 1)- invariant (Tran et al., 2025d).
We now verify the three metric axioms. First, for positive definiteness, is clear that STS-Soboley, μ)) 0 and STS-Sobolev,,(4 v) 2 0. If STS-Sobolev, (p, 1/) = 0, this implies O. Since the integrand is non-negative, this means Sp(uT, 1/) 0 for
S S "T) (T) almost all T E T. As S, is a metric, it follows that fTt = 1/ for almost all T. By the injectivity of the Spherical Radon Transform R conclude that the measures are equal, μ
Second, for symmetry, the property on each tree implies SpTT, = Sp(uT, NT). Therefore,
( ( do(
STS-Soboley, V,,
( Sp(UTI do())
-STS-Soboley,
(89)
Third, for the triangle inequality, the triangle inequality of S, for any 11, 12, 14 P(Sd) and then apply Minkowski's integral inequality:
(1, HATT dol do(7)² 1/2
STS So..Sooool,.,,//ee.e
(Sp(HL.T)H2TT) Sp(ea.T ( do(T)) 1/
S, 127) da))) T 1 + STS-Sobbley, (11, U22 + STT-Sobolle,
( (P do(
(90)
Thus, STS-Sobolev, is a metric on (())). For O + 1)-invariance, we aim show that for any E O(d + 11, STS- STTSSobbley,,(1, v) STS STS-SSboley,
Since is O(d + 1) 1)-invariant, the action of is an isometry, i.e., S, PT) Sp((tt)) (gTv)aT) (Tran et al., 2025d). Using this
35

---

Published as a conference paper at ICLR 2026
and change of variables, compute:
STS-Soboley, gtv)²*
(o
Sp(Cotn) (ot da(gT)
{ Sp(uT, uT PT) (T = (TS-Soboley, v)P.
(91)
Taking the p-th root of both sides, we conclude that STS-Soboley, is O(d + 1)-invariant.
Proof for Theorem E.3
Proof. We prove the theorem in two parts. First, we establish the equality for the case p = 1, and second, we prove the general incquality for any p E [1, oo).
Part We first prove equality for p = By definition, for the case p = 1, the Spherical Tree-Sliced Sobolev IPM is:
STS- SSs-Sobolev11 (H, - S (T)
(92)
The Spherical Tree-Sliced Wasserstein distance is defined as STSW D) =) fo Wi (HT, PT) (7). To prove the theorem, is sufficient to show the integrands are equal. We analyze the discrete form of the Sobolev IPM for = 1:
SITT UT) Σ Be 1 m(Te)l
-
(93)
For 1, the coefficient 3, simplifies Be = We. Substituting this result gives:
$I(HT, uT) = Σ w. Iu(7e) v(7)).
(94)
which is the known closed-form solution for the I-Wasserstein distance on tree, WI(HT, PT). Since the integrands are equal, their expectations are equal, proving that STS-Sobolevi (14, 2)) STSW(H,
Part Next, we prove the general inequality by showing that on any given tree T, the integrand is bounded as Sp(UT,, "T S WI(T,
This relies on two facts established in Appendix D.2: (1) the Sobolev coefficient 3e S for all
p 2 1, and (2) for probability measures, (p(re) - v(>e]] [o, 1]. which implies !... PP S |... for
p 2 1. Using these facts, we can bound the p-th power of the Sobolev IPM on tree:
Sp uT)
Σ
Be (u((e) v(76))
tEB
Σ
We (a(T)) - w(To)?
(since Br S wr)
Σ
We (u(re) - w(7e)l
(since (...I''......
cE
=WI(HTVV).
(95)
Integrating the inequality Sp(uT, "T) WI(HT,VT) over all trees T T directly yields the theorem and completes the proof.
Proof for Theorem E.4
Proof. The Monte
Carlo estimator
STS-Sobolev is STS-Soboley (14, v) = for
( Et Sp uT))) Let the random variable X; Sp(UT, VT, VT,,, where each Ti ~ a. The
36

---

Published as a conference paper at ICLR 2026
expected value of X, is μax [[X,] STS- (IS-Soboley,(H,))) Let the variance of X2 be finite, of
[ - V[X.]. By the Central Limit Theorem, the sample mean X Σ X; has variance
Our estimator is the functtin g(X) = X1/2. Applying the Delta Method, the variance of the estimator can be approximated by V[g(X)] s3 ((") where the derivative is g'(y) yy The Root Mean Squared Error (RMSE) is the square root the variance:
RMSE / Soboee VT.
3 Px
ox.
(96)
Since дa and ox finite constants independent of the number of samples I,, the Monte Carlo approximation error decays at the standard rate of (L--1/))
F EXPERIMENTAL DETAILS
F.1
RUNTIME AND MEMORY ANALYSIS
Table 6: Complexity Analysis of TS-Sobolev and STS-Sobolev.
Distamce
Opuration
Decriptton
Computation
Memory
Sampling
Random sampling
O(Lid)
D'Aid)
Projection
multiplicatoz pounts Dasance calculation and softmax
OLLAad
nd)
TS T5-Sobolev
Distance besed weight
O(bkmd)
OLEAn md)
spliting
Seeting
Sotting projected oneelimetss
Of legge)lgn))
O(
Coefiisient computation
Compuling csefiitients (AL)
O(Cbs)
Total
OLLAnd Lim
D2 Tesl)
Sampling
sappling spherical Manis multiplcation points and source Disance calcalation and softmax
O(ZAi)
D(2bd)
Projection
O(Led)
DiLd nd)
STS-Soboler
Distance based weight
O(Lknd)
Of0 Lid a md)
spliting
Sorting
Sorting projected cxoetinates
OfTa lngn) O(lka)
Offa)
Coeffiiient computation
Cumputing
O(1bs)
Total
O(2lmd Inllgg)
OfE md]
400
200
200
300
506
1000
Memory usage (MB)
40000
500
1000
30000
200
150-
100
10000
50
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
Number of supports (n)
Number of supports (n)
Figure 4: Execution time and memory usage of TS-Sobolev
Computational and Memory Complexity. We summarize the complexity of our proposed dis- tance measures in Table The overall computational and memory costs for TS TS-Sobolev and STS- Sobolev are identical to their respective counterparts, Db-TSW (Tran et al., 2025a) and STSW (Tran al., 2025d). This is because the additional step of computing the coefficients, as defined Equa- tion Equation (6), has low complexity of (LLn)) which is subsumed by the dominant terms of the projection and sorting operations.
Empirical Scaling Analysis. To verify our theoretical complexity, we benchmark the runtime and memory scaling of SS-Sobolevv.. and TS- Soboloele with respect to the number of support points (n) and the data dimension (d). For these experiments, fix the hyperparameters at L 2500 trees and & = 4 lines per tree, and run all tests on a single NVIDIA H100 GPU. Results are averaged over 10 runs.
37

---

Published as a conference paper at ICLR 2026
400
50000
300
1000
Memory usage (MB)
1000
250
30000
Runtima
200
150
10000
50
2000
3000
supports
supports
Figure 5: Execution time and memory usage of TS-Sobolev..
TSW SL
40
De TSN
-obslee
15-5nbeleg-3
(m/g)
20
10
Number Supports
Figure 6: Runtime Comparison of TS-Sobolev and other methods.
As detailed in Figures and 5, the results confirm our analysis. The left panel of each figure shows that runtime exhibits clear linear scaling with both n and d. Similarly, the right panel of each figure shows that memory usage also scales linearly. These empirical findings are fully consistent with the theoretical complexities presented earlier.
Runtime Comparison with Other Methods. We empirically compare the runtime of our method against its counterparts in Figure 6. For the tree-sliced methods (Db-TSW and TS- TS-Sobolev), we use L 2500 trees and k 4 lines per tree, while for Sliced Wasserstein (SW), use 10, 000 projections. All results are averaged over 10 independent runs. The experiment confirms that TS- Sobolev's runtime is nearly identical that of Db-TSW, the additional coefficient computation step introduces negligible overhead.
Runtime Across Order p. As also shown in Figure 6, the runtime of TSSSSboley remains sistent across different orders of p. This is a crucial advantage of our framework, demonstrating higher-order metries without incurring any performance penalty that it provides the flexibility TSW. compared to the standard
F.2 GRADIENT FLOW ON EUCLIDEAN SPACE AND ON THE SPHERE
Euclidean Datasets. Table reports the performance of our proposed methods compared with several baselines on the 8 Gaussians and Gaussian 30d datasets. We use L = 25 trees and k lines for tree-sliced methods and L = 100 projections for other sliced methods. We train for 2500 steps using the Adam optimizer with global learning rate of 0.005. For the TSW variant, use learning rate of 0.005 for the 25 Gaussians dataset and 0.05 for the Gaussian 30d dataset, as in Tran al. (2025a). For TSW-SL and Db-TSW,
p=1. Each distribution has 500 samples.
38

---

Published as a conference paper at ICLR 2026
Table Average I-Wasserstein distance (Wi) on Gaussian 30d (Euclidean)
Method
Iter 500
Iter 1000
Iter 1500
Iter 2000
Iter 2500
SW
23.5
23.2
22.9
22.6
22.3
SWGG
22.6
22.7
22.7
22.7
22.7
LCVSW
23.1
22.5
22.0
21.4
20.8
TSW-SL
21.4
20.6
20.0
19.4
18.8
Db-TSW
21.1
19.8
18.6
17.3
16.1
TS-Sobolev1.2
21.0
19.5
18.0
16.4
14.9
TS TS-Sobolev1.5
20.4
18.1
14.7
10.5
6.69
TS-Sobolev2
22.1
21.2
20.1
19.2
18.6
Table 8: Average Log I-Wasserstein distance (W1) on Mixture vMFs (Spherical)
Method
Epoch 50
Epoch 100
Epoch 150
Epoch 200
Epoch 250
SSW
-1.477
-1.740
-1.870
-1.951
-2.008
S3W
-1.213
-1.373
-1.451
-1.486
-1.509
RI-S3W (1)
-1.225
-1.479
-1.602
-1.672
-1.722
RI-S3W (5)
-1.423
-1.662
-1.782
-1.863
-1.913
ARI-S3W
-1.527
-1.792
-1.950
-2.057
-2.136
STSW
-1.530
-1.812
-1.969
-2.46
-2.082
STS Sobolevi.5
-1.842
-2.051
-2.097
-2.151
-2.139
STS- TTS-Sobolev2
-1.814
-2.054
-2.120
-2.150
-2.158
Spherical Datasets. The probability density of the von Mises-Fisher distribution with mean direc- tion μ E Sd is expressed as:
f(x; pz, 1s) = Ca(R:) exp(1<p
d/
where k > 0 controls concentration, and Ca(R) = (2T))//2 1/p0//11////1 1
is the normalization factor.
Following Bonet et al. (2022); Tran et al. (2024c; 2025d), we consider target distribution of 12 vMFs with 2400 samples (200 per vFM) where x 50 and
f1 = (-1, o, 0), /5 (0, 1),
μr = (1, o, 0), A4 (0, 1,
μa
-1,
μ4 = (1, -0.66.0
#6 (0,-1,-0) Mn1 = =0-0,0, -1,,
μs (0, -0)
(Ø, 0, -1); 1410 = (c) (,.0.1).
A12 = (-
1 + 55
where o-
We fix L = 200 trees and k = 5 lines for tree-sliced distance while using L = 1000 projections for the rest. ARI-S3W (30) has 30 rotations with a pool size of 1000. RI-S3W (1) and RI-S3W (5) have and rotations, respectively. All methods are trained using Adam (Kinga et al., 2015) optimizer with 1r 0.01 over 250 epochs. For STS Sobolev, we use learning rate of 0.05.
Evaluation using W1. To verify our performance gains, we re-evaluated the gradient flow ex- periments using the 1-Wasserstein (W,) distance. As detailed Table 7, the proposed method demonstrates significant improvements in the Euclidean setting under W, evaluation; notably, on the high-dimensional Gaussian 30d dataset, TS-Sobolev (p 1.5) achieves final Wi distance of 6.69, substantially outperforming the strongest bascline, Db-TSW (16.1). Furthermore, in the spher- ical setting (Table 8), STS-Soboler (p 2) achieves the lowest final log W1 distance of -2.158, surpassing both STSW (-2.082) and the strongest sliced baseline, ARI-S3W (-2.136).
Ablating ProjeProoet Settings. We conduct an ablation study to verify that the performance gains of the TS- Sobolev framework translate other tree projection settings. Theoretically. the core advantage our proposed method derived from the metric formulation itself and should therefore persist regardless of the specific tree structure splitting map employed. To empirically validate
39

---

Published as a conference paper at ICLR 2026
this, conducted additional Gradient Flow experiments the Gaussian 30d dataset, evaluating the method across four distinct configurations: combinations Chain versus Concurrent tree sampling strategies, and Uniform versus Distance-based splitting maps.
As presented in Table 9, 1SSoooolev1.2 consistently outperforms the standard TSW1 bascline across all four settings. This consistency confirms that the performance gains intrinsic the regularized Sobolev metric formulation rather than being specific to the default projection settings.
Table 9: Ablation study the Gaussian 30d dataset. We report the average Wasserstein distance (multiplied by 10-1) between source and target distributions at iteration 2500 across different tree structures and splitting maps. Results averaged over 10 runs.
Tree Sampling
Splitting Map
TSW1
TS-Sobolev
TS- TS-Sobolevi.6
TS-Sobolev2
Chain
Uniform
2.01
1.89
4.54
12.30
Chain
Distance
1.56
1.49
1.88
5.30
Concurrent
Uniform
1.93
1.83
3.37
11.10
Concurrent
Distance
1.78
1.40
1.51
3.68
F.3 DIFFUSION MODELS
Diffusion Models. Diffusion models (Sohl-Dickstein et al., 2015; Ho al., 2020) are a class of generative models renowned for producing high-quality samples. Their methodology is based a dual-process framework. The first fixed forward process, where data zo is progressively corrupted over 7' timesteps by adding Gaussian noise according predefined variance schedule, Sr. This noising cascade is defined by the transition kernel:
(x = Na
3,
The second a learned reverse process, where a neural network, parameterized by 0, is trained to reverse the corruption. At each timestep t, the model learns to predict the denoised sample 2,-1 from the noisy input It. This learned denoising step is also modeled as a Gaussian distribution:
(0 12:)) = (= (o( 1),
The model is trained by optimizing the Evidence Lower Bound (ELBO), which is equivalent minimizing the Kullback- (KL) divergence between the model's predicted distribution po(xe-1la7) and the true posterior q(
Denoising Diffusion GANs. A significant drawback of traditional diffusion models is their slow sampling speed, which stems from the large number of sequential steps (7) required. Denoising Diffusion GANs (DDGANs) (Xiao et al., 2021) address this inefficiency by reformulating the reverse process. Instead simple denoising network, DDGANs employ a conditional Generative Adver- sarial Network (GAN) for each reverse step. This approach allows for much larger and more expres- sive denoising transitions, drastically reducing the number of sampling steps needed- sometimes to enabling over 2000x speedups without substantial loss in sample quality.
few
The training objective for DDGANs has also evolved. While the original work relied on a standard adversarial loss, Nguyen al. (2024) successfully replaced with the Augmented Generalized Mini- batch Energy (AGME) distance. The AGME is sophisticated metric derived from the Generalized Mini-batch Energy (GME) distance (Salimans al., 2018), which quantifies the difference between two distributions by comparing the distances between mini-batches of their samples. The GME distance is defined as:
GME (1, v) 2E Py) -
where Px, Py are empirical measures from mani-batches and D is chosen base metric. The effec- tiveness of this training scheme highly dependent the choice of D. In our work, explore the performance of Sliced Wasserstein (SW) and our proposed Tree-Sliced Wasserstein (TSW) variants the base metric D within this framework.
40

---

Published as a conference paper at ICLR 2026
Implementation Details. Our experimental configuration closely follows the setup of Nguyen et al. (2024) and Tran et al. (2025a) for model architecture and core hyperparameters. All models trained for 1800 epochs.
For the tree-sliced methods, we set the number of sampled trees to = 2500 and lines per tree to k = 4, with sampling standard deviation of 0.1, per (Tran et al., 2025a). In contrast, for Sliced Wasserstein (SW) methods, we use = 10000 projections, consistent with (Nguyen et al,, 2024). We adopt the learning rates from the same work, setting them to Ira 1.25 10 4 for the discriminator and b 1.6 x 10-4 for the generator. All runtime evaluations are conducted with a batch size of 128 on two NVIDIA H100 GPUs. Our results for TS- TS-Sobolev are averaged over 10 runs while other results are obtained from previous results.
F.4 SULERIUPERVISED LEARNING
Encoder. In line with (Bonet et al., 2022: Tran et al., 2024c; 2025d), we train a ResNet18 (He et al., 2016) on CIFAR-10 10 data for 200 epochs with batch size of 512. Training uses SGD with Ir 0.05, momentum 0.9, and weight decay of 10-3 Data augmentations for creating positive pairs are aligned with earlier studies (Wang & Isola, 2020; Bonet et al., 2022; Tran et al., 2024c), including resizing, cropping, horizontal flipping, color jittering, and random grayscale conversion.
For tree-sliced methods, we use L = 200 trees, k = 20 lines for STSW, and À 10. We set = 200 projections for all other sliced distances. We report result in Table 3 where =
Linear Classifier. We then train linear classifier on the frozen representations produced by the encoder. Similar to prior works Bonet et al. (2022), training runs for 100 epochs using the Adam (Kinga et al., 2015) optimizer using a learning rate of 10-3, a weight decay of 0)2 at epochs 60 and 80.
F.5 TOPIC MODELING.
Topic modeling task (Blei et al., 2003) seeks to automatically extract distinct themes from collec- tions of text documents, revealing the underlying structure of a corpus. Recent approaches typically employ variational autoencoder (VAE) setup. in which the optimization balances accurate docu- ment reconstruction with a regularrzzttion that encourages the inferred topic distributions to resemble chosen prior (Srivastava & Sutton, 2017). Inspired by Nan al. (2019); Adhya & Sanyal (2025), we propose replacing the conventional KL--divergecce regularizer with a Wasserstein- -aased alterna- tive. This leads to the following objective:
inf Ep [CE(x, x)] + À ATS 62.97
(()).
where CE represents the cross- cross-entropy between the input document x (in ong-of-wortd representa- tion) and its reconstruction x. The variational posterior 9. (0)x x) is generated by encoder p, and the decoder &b maps topic mixtures 0 back to word distributions to form x.
Datasets. We evaluate our proposed methods on three well-known benchmark corpora used exten- sively for topic modeling research:
BBC (Greene & Cunningham, 2006): Comprising more than 2,000 news articles published by the BBC, grouped into 5 topical classes.
M10 (Pan et al., 2016): Extracted from the CiteSeer digital library, containing over 8,000 academic papers spanning 10 distinct research fields.
Preprocessing include lowercasing, punctuation removal, lemmatization, filtering out words shorter than three characters, and exclusion of documents with fewer than three words. Comprehensive statistics on these datasets after preprocessing are summarized in Table 10.
Evaluation Metries. To quantitatively measure model effectiveness, we consider topic coherence and diversity. Topic coherence is measured using the Cv (CV) metric (Röder et al., 2015), which correlates well with human interpretability, while topic diversity is assessed via the IRBO metric (Terragni et al., 2021), which evaluates how distinct the topics are. Topic coherence reflects the tent to which bigh-brbbbbilityy words within topics co-occur in documents, whereas diversity reflects how thematically different the topics from another.
41

---

Published as a conference paper at ICLR 2026
Table 10: Statistics and hyperparameters for datasets used.
Dataset statisties
Hyperparameters
Dataset
#Docs
#Labels
#Words
#Projections
Batch size
Dropout rate
Spherical prior
M10 BBC
8.355
10
1,696
2,000
64
0.5
Uniform vMF
2,225
5
2,949
8,000
256
0.05
Table 11: Topic diversity scores measured by IRBO (1) the BBC and M10 datasets
Method
BBC
M10
LDA (Blei et al., 2003)
0.917 = 0.016
0.915 + 0.012
ProdIDA (Srivastava & Sutton, 2017)
0.000
0.997 0.002
WTM (Nan et al., 2019)
0.997 0.003
0.812 1 0.060
Euclidean setting
SW-TM (Bonneel al., 2015)
0.998 = 0.004
0.971 = 0.008
RPSW-TM (Nguyen et al., 2024)
0.904 = 0.010
0.969 1 0.010
EBRPSW-TM (Nguyen et al., 2024)
0.999 = 0.003
0.963 + 0.013
TSW-SL-TM (Tran et al., 2024d)
0.999 = 0.003
0,981 + 0.005
Db-TSW-TM (Tran
2025a)
1.000
0.805 0.044
TS 7S-Soholev--TM (ours)
1.000 0.000
0.984 + 0.006
Spherical setting
S2WTM (Adhya & Sanyal, 2025; Bonet et al., 2022)
0.995 = 0.008
0.944 + 0.006
S3W-TM (Tran et al., 2024c)
0.987 = 0.014
0.763 + 0.079
LSSOT-TM (Liu al., 2025)
0.994 = 0.004
0.922 - 0.026
STSW-TM (Tran et al., 2025d)
0.993 = 0.005
0.854 1 0.028
STS 1155-obollyy--Th (ours)
0.996 0.005
0.895 0.027
Training Protocol. The experiments are conducted using the OCTIS framework (Terragni et al., 2021), adhering to the setup described in (Adhya & Sanyal, 2025). Each model is trained for 100 epochs, employing a Dirichlet prior when operating in the Fuclidean latent space, while parameters for the spherical latent space prior specified in Table 10. The regularization weight À is system- atically varied between 0.5 and 10 in steps of 0.5. For approaches involving tree-based objectives, the number of trees is fixed 100. Additional training configurations can be found in Table 10.
Topic Diversity. The topic diversity results measured by IRBO scores 1 on the DBLP, MI0, and BBC datasets are displayed in Table 11. Notably, our proposed methods achieved comparable topic diversity and superior topic coherence relative to the baselines, underscoring their practical advan- tages.
F.6 EFFECTS OF THE NUMBER OF TREES (T.) AND LINES PER TREE (k-)
The computational complexity of the framework scales linearly with both the number of sampled L and the number of lines per tree k, making the selection of these hyperparameters a key practical consideration.
Number of Trees (L). The parameter L directly governs the precision of the Monte Carlo estimate. Since the approximation error decays at a rate of (L--1/2) sufficiently large L is necessary to ensure accuracy. Our experiments confirm this, showing that higher values of L consistently lead to improved performance on downstream tasks.
Lines per Tree (k). The parameter & controls the geometric expressiveness of each individual tree structure, We observe empirically that using multiple lines (k: > 1) yields significantly better results than the h = 1 case (which simplifies to standard sliced distance). This confirms that the richer geometry of the tree structure is vital for performance.
The off between L and k. While a larger l can intuitively capture more intricate data ge- ometries, it also dramatically expands the space of possible tree structures. This, in turn, may require
42

---

Published as a conference paper at ICLR 2026
tog Aissessein Ustane Coggnnsars
Masserstss Jotanc (gaussin (M Small
-03
-33
-25
Relutions
Rerosions
Figure 7: Loog Wasserstein Distance between two distributions for (1.2, 1.5, 2,6,
a much larger number of samples I. to ensure the space of trees is adequately represented. Our ex- periments indicate that naively increasing & to very high values does not always improve results, likely due to this sampling challenge. Developing methods that can leverage the expressiveness of high-k trees without requiring a prohibitively large L remains important direction for future research.
F.7 EFFECTS OF THE ORDER (p)
We evaluate the sensitivity of our method to the hyperparameter p gradient flow task with the 25 Gaussians and Gaussian 20d datasets. For these experiments, we use L = 25 trees and l: = 4 lines, with learning rates set to 0.005 and 0.05, respectively. Figure 7 plots the log W2 distance between the source and target distributions over 1000 steps for several choices of
The results reveal clear trade-off. Very high values of p lead unstable training and diver- gence, as the associated gradients can become excessively large. Conversely, lower values of generally lead to more stable training dynamics and consistently strong performance. However, moderately high p can also be beneficial, as the increased variance in the training dynamics can help the model escape local optima. This is supported by our main experiments Section where larger sometimes yields superior results. Based on this analysis, we recommend selecting within the range [1.0, 2.0] as robust starting point for tuning.
Theoretical Analysis. Values of 2 introduce significant optimization challenges due to the extreme scaling of the structural weights. The weighting term w(r))1-> contains a negative exponent that grows in magnitude as p increases. This creates loss landscape where weights can shift dras- tically based on the subtree volume: weights for edges near the root can vanish entirely (causing a loss of global structural guidance), while weights for edges with small volumes can grow dispropor- tionately large (if the effective weight base is small), leading to exploding gradients. This numerical instability makes the optimization difficult to tune and prone divergence.
The specific success of = 2 likely stems from its unique geometric and computational balance. Geometrically, gradients are linear with respect to the error (12-1 A). providing the most con- sistent and stable signal across different scales of error. Computationallyy, where the edge coefficient 3e takes logarithmic form log(1 I+w(,.) ). in contrast the polyno- mial form for other values. This logarithmic scaling naturally compresses the dynamic range of the spatial weights. It prevents the bias against global structure from becoming too extreme, avoiding the vanishing weight problem, and allows p to effectively prioritize fine-grained local details (leaves) while retaining sufficient sensitivity global alignment (root).
= is a special case
In this paper, we prioritize gradient-based optimization, which the most critical application of sliced distances. As detailed in F.8, setting p > 1 generally yields smoother gradients, making it more suitable for optimization tasks. Conversely, for = 1, the weighting function emphasizes global discrepancies, and the transport cost follows the L, norm. These properties render the p setting robust to local outliers and effective for estimating global changes, which can be particularly beneficial for tasks involving noisy data.
43

---

Published as a conference paper at ICLR 2026
F.8 ANALYSIS OF ADVANTAGES OF HIGHER ORDER (p 1)
TS-Sobolev is motivated by the need for a computationally efficient metric of order > 1. While the standard Tree-Sliced Wasserstein distance lacks a tractable solution for orders p > these higher orders is important for gradient-based learning, as y-Wasserstein metrics offer smoother gradients compared to the p 1 case (Peyré et al., 2019). Beyond simply enabling tractability, our analysis re- veals that TS-Sobolcv (p > 1) introduces distinct advantages over standard p- Wasserstein. Overall, the better performance of TS-Sobolev stems from two complementary mechanisms: the improved optimization landscape inherent to the L cost function and the preservation of fine-grained features introduced by the weighting function.
Improved Optimization Landscape. The choice of p dictates the convexity and smoothness of
the underlying optimization objective. Let ((A.) == |A.|" = [m( (u(::))-((>)) P. as in Equation (5), denote the unweighted transport loss associated with the mass discrepancy Ae on given edge. For p = 1, the Tree-Sliced Wasserstein (TSW) distance behaves analogously to an L' loss. This formulation lacks strict convexity, which implies that optimization problems solving for measures or / may admit non-unique minimizers, leading to potential instability (Santambrogio, 2015; Villani, 2003). Furthermore, the gradient magnitude for p 1 remains constant (1 (VL) x i., proportional to a constant) regardless of the proximity to the target, often leading to oscillations around the optimum unless the learning rate is carefully annealed. In contrast, for p > 1, the cost function becomes strictly convex, ensuring unique geodesics and well-conditioned gradient signals (Peyré et al.. 2019). Crucially, the gradient magnitude scales with the transport cost, following [ x |A.|-- This property ensures that gradients are large when distributions are distinct and vanish smoothly as A. 0, facilitating stable fine-tuning and convergence.
These better optimization characteristics directly translate into improved performance in the Gradi- ent Flow experiments across both Euclidean and Spherical benchmarks in Section 4.
Preserving Fine-Grained Structure. Beyond the optimization benefits inherent to standard Wp metries, TooSolllev employs a weighting mechanism that uniquely prioritizes the preservation of fine-grained features. As derived in Equation 4, the metric minimizes a cost weighted by the term PP) where (xx) = + (A(x)) represents the weight of the subtree rooted at æ. For p > 1, this weighting factor decays the subtree size w w((x)) increases. Consequently, this mechanism downscale dominant gradients arising from the root and upper levels of the tree (where si (A(x)) is large), preventing global mass shifts from overwhelming the optimizer. Conversely, nodes deeper in the tree (near the leaves) possess smaller subtree weights, resulting in substantially larger values for (xxx)>_P This effectively concentrates the optimization signal on minimizing local discrepancies. ensuring the capture and preservation of fine- grained details.
In image generation. fine-grained details correspond to textures and sharp edges. Consequently, this capacity to prioritize local feature details is a key factor driving the enhanced sample quality and sharpness observed in our large-scale diffusion training.
features such as intricate
To further demonstrate the ability to capture fine-grained high-frequency signals, we conducted a controlled experiment using a ID synthetic signal composed of distinct frequency modes. We defined target probability density p(z) on the domain [0, 1] as mixture of sine waves:
p(x) x 1 + 0.5 sin (2m klow 22 + 0.3 sin(2 -highh : x))
(97)
where klow = 2 represents low-frequency signals and khigh 20 represents high-frequent signals. We initialized N = 10000 particles from a uniform distribution [4 1] and optimized their positions via gradient flow to minimize the distance to p. comparing TSW against TS-Sobolevz. To quantify the capture of frequency modes, we computed the Discrete Fourier Transform (DFT), denoted as F, of both the particle density e and the target p. We calculated the relative spectral error as the magnitude of the difference between their spectral components: now-frequeccy error [btow] F(p)]kiow) and high-frreqencyy error [F()) [Knigh] - [kkglll
The quantitative results are summarized in Table 12, and the error trajectories are visualized in Figure 8. In the low-frequency regime, both metries perform similarly. At final, TSW achieves low-frequency error of 34.68 compared to 33.94 for Sobolev, indicating both effectively capture global structure. However, significant disparity emerges in the fighffreuuency regime. While the high- high-frequency error for TSW plateaus, the error for TS TS-Sobolev, consistently decreases, reaching 14. compared to 20.15 for TSW. This clear trend demonstrates that TS "os-ooolee, captures high-
44

---

Published as a conference paper at ICLR 2026
Table 12: eey-Frequyncyy and Hiqu-Erry Error (1) over 10 runs.
Low Frequency Error TSW
High Frequency Error TSW
Iter
TS-Sobolev2
TS-Sobolev2
250
48.18 + 1.45
47.80 1.27
25. 24 2.06
25.29 = 2.55
500
41.55 + 2.12
43.98 1.80
21.01 + 1.94
20.45 + 2.18
750
39.96
+ 2.46
39.25
2.24
19.24 + 0.90
16.05 + 1.50
1000
34.68
+ 2.75
33.94
+ 2.55
20.15 + 1.39
14.16 + 1.15
Low Frequency Error Error
High Frequency Error
TSV
TSV
T5 Socoice
T5 Soso
45.0
37.5
15.0
200
500 1000
200
Figure 8: The evolution of error for low-frequency (A: = 2) and high-frequency (k: = 20) modes during gradient flow. While both methods reduce low-frequency error similar rate, TS-Sobolev (Orange) converges significantly better on the high-frequency component than TSW (Blue).
frequency signals significantly better than TSW, confirming our theoretical analysis regarding the preservation of fine-graines structure.
F.9 HARDWARE SETTINGS
All experiments were conducted on single NVIDIA A100 (40GB) GPU, with the exception of the denoising diffusion experiments, which were executed in parallel across NVIDIA H100 GPUs.
G BROADER IMPACTS
The ability accurately and efficiently compute distances between complex probability distribu- tions is foundational challenge in many scientific and industrial fields. The TS-Sobolev frame- work presented in this work contributes new tool for this task, with potential for broad societal impact. Applications span data-driven science, where more robust distributional comparison could refine medical image analysis for diagnosties, and generative AI, where may enable the devel- opment of higher-fidelity models for creative content. A distinct advantage of our approach is its suitability for dynamic-support measures, a critical capability for real-time systems in domains like financial modeling, logistics, and environmental science. By advancing this fundamental computa- tional primitive the comparison measures our work can help foster progress across a range of applications reliant sophisticated data analysis.
45


Reviews:
[{"review_id": "8zHPPok5DF", "reviewer": "Reviewer_UFk1", "summary": "### **Summary**\n\nThis paper introduces the **Tree-Sliced Sobolev Integral Probability Metric (TS-Sobolev)** and its **spherical variant (STS-Sobolev)** — new metrics for comparing probability distributions that generalize the Tree-Sliced Wasserstein (TSW) distance to any order $p \\ge 1$.  \n\nThe main idea is to replace the 1-Wasserstein distance on tree metric spaces (used in TSW) with a **regularized Sobolev Integral Probability Metric (IPM)** that admits a **closed-form solution**. This substitution retains the computational efficiency of TSW (≈ $O(Lkn \\log n)$) while enabling higher-order Sobolev metrics.  \n\nThe authors establish key **theoretical guarantees**, including:\n- Metricity and $E(d)$-invariance,\n- A formal connection showing $\\text{TS-Sobolev}_{p=1} = \\text{TSW}$,\n- Monte Carlo convergence rate $O(L^{-1/2})$.\n\nThe proposed methods are evaluated on a wide range of **Euclidean and spherical tasks**, including gradient flows, diffusion-based generative modeling (DDGAN on CIFAR-10), self-supervised learning on the sphere, and topic modeling on BBC and M10 datasets.  \n\nAcross all benchmarks, TS-Sobolev and STS-Sobolev consistently **outperform Sliced and Tree-Sliced Wasserstein baselines** in accuracy or convergence rate, while maintaining similar runtime. The paper concludes that Sobolev-based slicing provides a scalable and flexible alternative to classical TSW for both Euclidean and non-Euclidean domains.", "questions": "**Q1.**  \nSuppose the tree $\\mathbb{T} \\subset \\mathbb{R}$, i.e., all nodes lie on a straight line. In this case, the tree-sliced construction should coincide with the original sliced OT construction. Under this 1D setting, can we claim that the regularized Sobolev IPM (Eq. (4)) or the Sobolev IPM (Eq. (3)) is equivalent to $W_p(\\mu,\\nu)$?  \nMy understanding is that in 1D Eq. (4) reduces to  \n$$\n\\int_{\\mathbb{R}} \\big|\\mu((-\\infty,x]) - \\nu((-\\infty,x])\\big|^p \\, dx,\n$$  \nwhich differs from the standard 1D OT definition  \n$$\nW_p^p(\\mu,\\nu) = \\int_0^1 |F_\\mu^{-1}(t) - F_\\nu^{-1}(t)|^p \\, dt.\n$$  \nIf this understanding is correct, then the proposed Tree-Sliced Sobolev IPM (Def. 3.1) cannot recover the classical sliced OT formulation  \n$$\n\\int_{\\mathbb{S}^{d-1}} W_p^p((P_\\theta)_\\# \\mu, (P_\\theta)_\\# \\nu) \\, d\\sigma(\\theta),\n$$  \nsince $S_p^p(\\mu,\\nu)$ and $W_p^p(\\mu,\\nu)$ are not equivalent in 1D. Could the authors clarify this or specify under what conditions the two coincide?  \n\n**Q2.**  \nThe paper motivates Sobolev IPM as more “frequency-aware” and capable of mitigating spectral bias, but the experiments mainly show performance improvements without spectral evidence. Could the authors include a more direct spectral analysis (e.g., using synthetic 1D or 2D signals with known low/high-frequency components) to verify that the proposed tree-sliced Sobolev formulation indeed captures or preserves high-frequency modes better than standard sliced OT?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "### **Strengths**\n\n#### **Novel Theoretical Contribution**\n- Introduces a closed-form **Sobolev IPM** on trees that generalizes **1-Wasserstein** while retaining efficiency.  \n- Provides formal proofs of **metricity**, **$E(d)$-invariance**, and **Monte Carlo convergence rate**.  \n- Connects **Sobolev IPM** and **tree-sliced methods** in a clean, unified framework.\n\n#### **Scalability**\n- Maintains the computational complexity $O(Lkn \\log n)$ of existing TSW methods, enabling practical scalability to large datasets.  \n- Extensible to **spherical manifolds** via the **spherical Radon transform**.\n\n#### **Strong Empirical Results**\n- Consistently improves over both **SW** and **TSW** baselines across diverse domains.  \n- Demonstrates flexibility across geometries (**Euclidean** / **spherical**) and applications (gradient flows, SSL, generative modeling, topic modeling).\n\n#### **Clarity of Exposition**\n- Well-structured presentation: clear flow from background → formulation → theoretical properties → experiments.  \n- Figures and equations are informative, and notation aligns with prior **TSW literature**.", "weaknesses": "1. In the computation section, the authors claim a total complexity of $O(Lkn \\log n + LKd n)$, but it seems that the tree construction cost is not explicitly included. \n1.1 Moreover, the parameter $k$ (the number of lines per tree) controls a trade-off between computational efficiency and information preservation in $\\mathbb{R}^d$. A detailed discussion of this trade-off would strengthen the paper. In particular, it would be useful to know how large $k$ can grow before the overall method becomes slower than classical OT solvers such as Sinkhorn, whose complexity is $O(n^2 / \\epsilon)$. \n1.2 Clarifying whether $k$ depends on $n$ would also help to assess the true scalability of the approach.\n\n2. In the spherical setting, I suggest including additional sliced OT baselines for a more complete comparison — for example, “Spherical Sliced Optimal Transport” (arXiv:2411.06055) and “Stereographic Spherical Sliced Wasserstein Distances” (arXiv:2402.02345). These would provide a stronger empirical context and highlight the advantages of the proposed method under spherical geometry.\n\n3. The Sobolev IPM involves gradient norms $|\\nabla f|$, which are central to the metric’s definition. However, estimating $\\nabla f$ on discrete samples and through non-smooth tree-sliced projections may introduce numerical instability or high-variance gradients. Prior work on Sobolev-based metrics (Mroueh et al., ICLR 2018; Deshpande et al., NeurIPS 2019; Korotin et al., ICML 2021) has highlighted that $\\nabla f$ estimation can become unstable without smoothing or gradient regularization. The paper would benefit from analyzing or mitigating such potential instability in the proposed tree-sliced formulation.", "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": "hJI3BIlOXX", "reviewer": "Reviewer_ZHaq", "summary": "Tree-Sliced Wasserstein (TSW) practical use is constrained to the 1-Wasserstein because $W_p$ ($p>1$) for tree is costly. This paper proposes using regularized Sobolev IPM which has a computationally efficient closed-form solution on trees. THis recovers TSW exactly at $p = 1$ and is upper-bounded by TSW for $p > 1$, while maintaining the same $\\mathcal{O} (L k n log n)$ complexity.", "questions": "- Can the authors give more insight, perhaps related to the Sobolev IPM formulation involving derivatives, on why $p>1$ might lead to better optimization dynamics (e.g., smoother gradients, better conditioning) compared to $p=1$ (TSW), as suggested by the gradient flow results?\n- Could the authors please comment on the practical effect of the structural weighting $\\hat{w}(x) = 1 + \\omega(\\Lambda(x))$? Does this induce spatial sensitivity bias on the tree slices and how would it interact with the choice of $p$ for example?\n- How robust are the benefits of $p>1$ expected to be under different tree generation schemes (e.g., chain structures) or alternative splitting maps?\n- Appendix F.7 provides valuable empirical analysis showing $p \\in [1.0, 2.0]$ is effective and higher $p$ can be unstable . Can the authors offer any further heuristic or theoretical guidance? For instance, does the optimal $p$ appear correlated with data dimensionality or the nature of the task (e.g., generation vs. optimization)? Why might $p=2$ perform particularly well across several diverse tasks (Tables 2, 4, 5)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- A key achievement is generalizing to $p > 1$ without sacrificing computational efficiency. The closed-form expression for $\\hat{\\mathcal{S}}_p$ on discrete measures (Theorem B.6, Eq. 28) relies on edge coefficients $\\beta_e$ that depend on $p$ but are efficiently pre-computable, adding negligible $\\mathcal{O}(Lkn)$ overhead. Thus TS-Sobolev retains the dominant $O(Lkn \\log n)$ complexity of TSW for \\textit{any} $p \\ge 1$. This claim is validated by the empirical runtime analysis (Appendix F.1, Figs 4-6), showing near-identical wall-clock times regardless of $p$.\n- Good experimental results. For example,\n    - Achieves faster convergence and lower final $W_2$ error compared to TSW ($p=1$) and SW variants in Euclidean settings (Table 1). Similarly outperforms spherical TSW (STSW) in spherical gradient flows (Table 4).\n    - with DDGAN training gives sota FID scores on CIFAR-10, significantly improving upon TSW-based DDGANs with identical training time per epoch\n    - Strong performance in spherical SSL (Table 3) and achieves the highest topic coherence ($C_V$) in both Euclidean and Spherical topic modeling benchmarks (Table 5).", "weaknesses": "- The main justification for using $\\hat{\\mathcal{S}}_p$ is its tractability. However, the paper offers limited intuition on the gain or loss by using this specific Sobolev-based IPM instead of $W_p$ on trees for $p>1$. $\\hat{\\mathcal{S}}_p$ relates to the weighted $L^{p'}$ norm of the critic function's derivative on the tree. How does minimizing this discrepancy differ fundamentally from minimizing transport cost ($d(x,y)^p$)? Does $\\hat{\\mathcal{S}}_p$ emphasize smoothness or local variations differently than $W_p$? \n\n- The closed-form tractability of $\\hat{\\mathcal{S}}_p$ comes from a specific weighting function $\\hat{w}(x) = 1 + \\omega(\\Lambda(x))$ inherent in the underlying Sobolev norm equivalence on trees (Theorem B.3). This weight depends on the measure (length) of the subtree $\\Lambda(x)$ below point $x$ relative to the root (Eq. 22). It is not clear to the reviewer the geometric implications of this structural weighting. For example, does this make $\\hat{\\mathcal{S}}_p$ (and thus TS-Sobolev) inherently more sensitive to distributional differences occurring deeper in the tree projection (larger $\\omega(\\Lambda(x))$) versus closer to the root? \n\n- The experiments consistently use the concurrent-line tree structure [cite: lines 1097-1102] and softmax splitting map, following configurations from recent TSW paper. While sensible for direct comparison, this limits the exploration of how TS-Sobolev performs under different structural assumptions. Appendix F.6 analyzes the $L$ vs. $k$ trade-off but doesn't compare tree topologies (e.g., concurrent vs chain or splitting maps within the TS-Sobolev context. The observed benefit of $p>1$ might interact with these choices.", "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": "2FEZSr1hPT", "reviewer": "Reviewer_yezB", "summary": "In this work, the authors proposed a new type of tree-sliced metric to capture the $p$ -Wasserstein distance between two probability measures defined either in $\\mathbb{R}^d$ or on a hypersphere. The main motivation is that recent advances in tree-sliced Wasserstein (TSW) [1,2] only provide closed-form expressions for $p = 1$, which is quite limited. On the other hand, the regularized Sobolev Integral Probability Metric -- Sobolev IPM -- (for any order $p \\geq 1$) between two distributions on trees is known in closed form [3]. The idea is as follows: one randomly samples a tree structure (a slice) based on a procedure in [2], then computes the closed-form regularized Sobolev IPM (of order $p$) between two distributions within that tree structure. This process is repeated multiple times to obtain a Monte Carlo estimation of the proposed metric which is defined as the expectation of this MC estimator.\n\nThe authors showed that the proposed metric is indeed a valid metric (e.g., symmetric and satisfying the triangle inequality). When $p = 1$, it reduces to the tree-sliced Wasserstein metric, and for $p > 1$, it serves as a lower bound of TSW.  \n\nIn the experiments, the proposed metric was applied in several contexts: gradient flows in either Euclidean space or on a hypersphere, as part of the objective function in training diffusion models, and in topic modeling. In most cases, the proposed metric outperforms other Wasserstein-based metrics and demonstrates improvements in downstream tasks.  \n\n**References**\n\n[1] Tran, V. H., Pham, T., Tran, T., Le, T., & Nguyen, T. M. (2024). Tree-sliced Wasserstein distance on a system of lines. arXiv e-prints, arXiv-2406. \n\n[2] Tran, H. V., Nguyen, K. N., Pham, T., Chu, T. T., Le, T., & Nguyen, T. M. (2025). Distance-based tree-sliced Wasserstein distance. arXiv preprint arXiv:2503.11050. \n\n[3] Le, T., Nguyen, T., Hino, H., & Fukumizu, K. (2025). Scalable Sobolev IPM for Probability Measures on a Graph. arXiv preprint arXiv:2502.00737.", "questions": "- In 2D, when generating a system of lines, they inevitably intersect, allowing one to remove some intersection points to construct a tree. I wonder how this idea extends to higher dimensions ($d > 2$), where such lines generally do not intersect.  \n\n- Is the method applicable to more general Riemannian manifolds beyond the hypersphere? In principle, one could sample geodesics in those manifolds, so it would be interesting to see whether the approach can be generalized.\n\n- As mentioned before, it would be helpful to clarify why we really need the p-Wasserstein distance, and in what ways the 1-Wasserstein distance falls short.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The idea of combining tree-sliced Wasserstein with the Sobolev IPM to handle the general p-Wasserstein distance is novel and provides a practical way to approximate the $p$-Wasserstein distance in practice.  \n\n- The proposed approach introduces minimal computational overhead compared to existing Sliced or Tree-Sliced Wasserstein methods.  \n\n- The experiments are extensive and consistently demonstrate improved performance over previous works.", "weaknesses": "- The necessity of approximating the p-Wasserstein distance is not clearly motivated. It would be helpful to clarify what specific advantages the p-Wasserstein distance offers over the 1-Wasserstein distance, which the existing tree-sliced Wasserstein approach is already able to capture.  \n\n- Some experimental comparisons may favor the proposed method. For example, in the gradient flow experiment on the sphere (and possibly also in the Euclidean space), using the 2-Wasserstein distance as the evaluation criterion could naturally benefit the proposed method over Tree-Sliced Wasserstein [2], which is inherently based on the 1-Wasserstein distance. I wonder what the results look like if you used 1-Wasserstein distance as the criterion instead.\n\n- The writing could be improved for clarity and structure. For instance, the definition of the regularized Sobolev IPM should be presented in the main text, as it is a central component of the proposed method. Additionally, the transition from “Tree Metric Spaces” to “Tree-Sliced Wasserstein Distance” is somewhat abrupt and may confuse readers. For example, the paragraph on “Tree Metric Spaces” might lead one to believe that the space consists solely of nodes, while in “Tree-Sliced Wasserstein Distance,” the entire tree structure (including edges) is considered. Improving the connection between these sections and providing a smoother introduction would enhance readability.  \n\n**Reference**\n\n[2] Tran, H. V., Nguyen, K. N., Pham, T., Chu, T. T., Le, T., & Nguyen, T. M. (2025). *Distance-based tree-sliced Wasserstein distance.* arXiv preprint arXiv:2503.11050.", "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)