Tree-Sliced Variants of Wasserstein Distances

Tam Le, Makoto Yamada, Kenji Fukumizu, Marco Cuturi

Abstract: Optimal transport ($\OT$) theory defines a powerful set of tools to compare probability distributions. $\OT$ suffers however from a few drawbacks, computational and statistical, which have encouraged the proposal of several regularizations/simplifications in the recent literature, one of the most notable being the \textit{sliced} variant of OT which considers univariate projections. We consider in this work a particular family of ground metrics, namely \textit{tree metrics}, which yield negative definite $\OT$ metrics that can be computed in a closed form, of which the sliced-Wasserstein distance is a particular case (the tree is a chain). We propose a positive definite tree-Wasserstein ($\TW$) kernel building on this, and also propose two ways to build tree metrics in both low-dimensional and high-dimensional spaces. We propose the tree-sliced Wasserstein distance, using averages over random tree-metrics. We empirically illustrate that the proposed $\TW$ kernel compares favorably with other baselines on several benchmark datasets.

Code Link: https://github.com/lttam/TreeWasserstein

CMT Num: 6654