Dimensionality Reduction for Wasserstein Barycenter

TL;DR

This paper introduces a dimensionality reduction approach for Wasserstein barycenters that preserves solution quality while significantly reducing computational complexity, addressing the curse of dimensionality in high-dimensional probability distribution analysis.

Contribution

The paper develops randomized dimensionality reduction techniques for Wasserstein barycenters, providing optimal bounds and a coreset construction to improve efficiency and scalability.

Findings

Dimensionality reduction preserves barycenter cost with small error.

Methods are optimal up to constant factors.

Experimental results show significant speedup with maintained accuracy.

Abstract

The Wasserstein barycenter is a geometric construct which captures the notion of centrality among probability distributions, and which has found many applications in machine learning. However, most algorithms for finding even an approximate barycenter suffer an exponential dependence on the dimension $d$ of the underlying space of the distributions. In order to cope with this "curse of dimensionality," we study dimensionality reduction techniques for the Wasserstein barycenter problem. When the barycenter is restricted to support of size $n$, we show that randomized dimensionality reduction can be used to map the problem to a space of dimension $O(\log n)$ independent of both $d$ and $k$, and that \emph{any} solution found in the reduced dimension will have its cost preserved up to arbitrary small error in the original space. We provide matching upper and lower bounds on the size of the reduced dimension, showing that our methods are optimal up to constant factors. We also provide a coreset construction for the Wasserstein barycenter problem that significantly decreases the number of input distributions. The coresets can be used in conjunction with random projections and thus further improve computation time. Lastly, our experimental results validate the speedup provided by dimensionality reduction while maintaining solution quality.