Mixed-Curvature Tree-Sliced Wasserstein Distance

Published in The Fourteenth International Conference on Learning Representations (ICLR 2026), 2025

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