Tree-Sliced Wasserstein Distance with Nonlinear Projection

Thanh Tran, Hoang V. Tran, Thanh Chu, Huyen Trang Pham, Laurent Ghaoui, Tam Le, Tan Minh Nguyen

TL;DR: We propose a novel nonlinear projectional framework for the Tree-Sliced Wasserstein distance, substituting the linear projections in earlier versions with general projections

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.

Lay Summary: Tree-Sliced methods are a new way to compare probability distributions by using trees instead of lines, making it easier to capture complex shapes in the data while keeping computations fast. Building on this idea, we introduce a new version that uses more flexible (nonlinear) projections but still keeps the math valid and efficient. We design these projections for both flat (Euclidean) and curved (spherical) data, and test them in various tasks like learning representations and generating data. Our method shows better results than previous approaches.

Primary Area: Deep Learning

Keywords: tree-sliced wasserstein distance, optimal transport

Submission Number: 11529