Tree-sliced Sobolev IPM

ICLR 2026 Conference Submission5088 Authors

Keywords: sliced optimal transport, sobolev ipm, tree-sliced wasserstein distance, tree wasserstein distance

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\ge1$; 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.

Primary Area: learning on graphs and other geometries & topologies

Submission Number: 5088