Efficient Algorithms for Computing a Partial $p$-Wasserstein Barycenter in Metric Spaces

Given a collection of $k$ probability distributions over a total of $n$ points in a space $\Omega$ equipped with a pairwise distance function $\text{d}$, the task of computing an average probability distribution with respect to the optimal transport cost is known in the literature as the *Wasserstein barycenter problem*. However, like OT plans, Wasserstein barycenters are also sensitive to the presence of outliers in each distribution. The *$\alpha$-partial Wasserstein barycenter problem* for $\alpha \in (0,1]$, which addresses this issue, asks to compute a measure of $\alpha$ mass for which the average unbalanced OT cost from the input distributions to this measure is minimized. In this paper, we design provably efficient approximation algorithms for the $\alpha$-partial Wasserstein barycenter problem. Our first algorithm computes an approximate fixed-support $\alpha$-partial Wasserstein barycenter in arbitrary spaces in roughly $mn\cdot {\delta }^{-2}$ time, where $m$ is the size of the fixed support and $\delta$ is a small approximation error in the mass. The second algorithm improves the run time in the free-support setting for metric spaces by a factor of $m$ at the expense of an additional $ϵ$ approximation error in the cost of the OT plans. We additionally prove that an exact solution to the $\alpha$-partial Wasserstein barycenter problem can be computed in roughly $n^{O(d)}$ time.