Tree-Sliced Wasserstein Distance with Nonlinear Projection

TL;DR

This paper introduces a nonlinear projection framework for Tree-Sliced Wasserstein distance, improving its ability to capture topological structures in measures while maintaining computational efficiency, with applications in machine learning tasks.

Contribution

It proposes a novel nonlinear projectional approach for Tree-Sliced Wasserstein distance, ensuring injectivity and applicability to Euclidean and spherical measures, advancing the state-of-the-art in optimal transport metrics.

Findings

Enhanced metrics for Euclidean and spherical measures

Significant improvements in gradient flows, self-supervised learning, and generative models

Validated through extensive numerical experiments

Abstract

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