Tree-Sliced Wasserstein Distance: A Geometric Perspective

2024-06-19

AI Summary

This work addresses the topological information loss and degraded geometric awareness inherent in the Sliced Wasserstein (SW) distance due to its reliance on linear projections. To resolve this, we propose the Tree-Sliced Wasserstein distance with Splitting Maps (TSW-SL)—the first extension of SW to metric tree structures. Our core innovation lies in introducing a splitting map and a tree-domain Radon transform, rigorously proving their invertibility and metric properties while preserving the closed-form optimal transport solution. Theoretically, this generalizes both the Radon transform and optimal transport frameworks to non-Euclidean, hierarchical domains. Algorithmically, TSW-SL enables efficient gradient flow optimization via differentiable tree projections. Experiments demonstrate that TSW-SL consistently outperforms SW and its variants in gradient flow simulation, image style transfer, and generative modeling—achieving superior geometric robustness without sacrificing computational efficiency.

Technology Category

• Machine Learning: Learning with Manifolds
• Planning, Routing, and Scheduling: Optimization of Spatio-temporal Systems
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Many variants of Optimal Transport (OT) have been developed to address its heavy computation. Among them, notably, Sliced Wasserstein (SW) is widely used for application domains by projecting the OT problem onto one-dimensional lines, and leveraging the closed-form expression of the univariate OT to reduce the computational burden. However, projecting measures onto low-dimensional spaces can lead to a loss of topological information. To mitigate this issue, in this work, we propose to replace one-dimensional lines with a more intricate structure, called tree systems. This structure is metrizable by a tree metric, which yields a closed-form expression for OT problems on tree systems. We provide an extensive theoretical analysis to formally define tree systems with their topological properties, introduce the concept of splitting maps, which operate as the projection mechanism onto these structures, then finally propose a novel variant of Radon transform for tree systems and verify its injectivity. This framework leads to an efficient metric between measures, termed Tree-Sliced Wasserstein distance on Systems of Lines (TSW-SL). By conducting a variety of experiments on gradient flows, image style transfer, and generative models, we illustrate that our proposed approach performs favorably compared to SW and its variants.

Problem

• Addresses computational burden in Optimal Transport via tree systems
• Mitigates loss of topological info in low-dimensional projections
• Proposes Tree-Sliced Wasserstein distance for efficient metric evaluation

Innovation

• Uses tree systems instead of one-dimensional lines
• Introduces splitting maps for projection mechanism
• Proposes novel Radon transform for tree systems