Tree-Sliced Wasserstein Distance: A Geometric Perspective

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 \emph{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.

Optimal Transport (OT) is a powerful tool but often too slow to compute. A popular faster version, called Sliced Wasserstein (SW), simplifies the problem by projecting it onto many lines, which can miss important shape information. In this work, we propose a new method that uses a more complex structure called tree systems instead of simple lines. This lets us keep more of the original data’s structure while still computing OT efficiently. We introduce a new way to project onto these trees and show that it works well. Our method, called Tree-Sliced Wasserstein on Systems of Lines (TSW-SL), runs efficiently and outperforms SW in experiments like image style transfer and generative modeling.