Tree-Sliced Wasserstein Distance with Nonlinear Projection
This work provides an incremental improvement in optimal transport metrics for machine learning tasks such as self-supervised learning and generative modeling.
The authors tackled the problem of enhancing the Tree-Sliced Wasserstein distance by introducing a nonlinear projection framework to better capture topological structures while maintaining computational efficiency, resulting in significant improvements over existing methods in applications like gradient flows and generative models.
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.