Tree-Sliced Wasserstein Distance with Nonlinear Projection

2025-05-02

AI Summary

Computing distances between high-dimensional probability measures often faces a trade-off between structural fidelity and computational efficiency. Method: This paper proposes the Nonlinear Tree-Sliced Wasserstein (NTSW) distance, introducing the first invertible nonlinear projection framework satisfying Radon transform injectivity—replacing conventional linear projections—while preserving the low computational complexity (O(n log n)) of tree-sliced methods. Contribution/Results: Theoretically, NTSW yields a well-defined metric compatible with both Euclidean and spherical measures. Algorithmically, it integrates sliced optimal transport, gradient flow modeling, and self-supervised optimization. Experiments demonstrate that NTSW achieves 12–27% higher accuracy than state-of-the-art Sliced and Tree-Sliced Wasserstein methods in generative modeling, representation learning, and gradient flow simulation, with comparable computational cost.

Technology Category

• Search and Optimization: Non-convex Optimization
• Machine Learning: Learning with Manifolds
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web evaluation methodologies and metrics
• Security and Privacy: Large-scale security measurements

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.

Problem

• Enhancing Tree-Sliced Wasserstein distance with nonlinear projections
• Preserving injectivity and metric properties in Radon Transform
• Improving computational efficiency for Euclidean and spherical measures

Innovation

• Nonlinear projection for Tree-Sliced Wasserstein distance
• Injectivity preservation in Radon Transform
• Efficient metrics for Euclidean and spherical measures