Parallel Streaming Wasserstein Barycenters
Efficiently aggregating data from different sources is a challenging problem, particularly when samples from each source are distributed differently. These differences can be inherent to the inference task or present for other reasons: sensors in a sensor network may be placed far apart, affecting their individual measurements. Conversely, it is computationally advantageous to split Bayesian inference tasks across subsets of data, but data need not be identically distributed across subsets. One principled way to fuse probability distributions is via the lens of optimal transport: the Wasserstein barycenter is a single distribution that summarizes a collection of input measures while respecting their geometry. However, computing the barycenter scales poorly and requires discretization of all input distributions and the barycenter itself. Improving on this situation, we present a scalable, communication-efficient, parallel algorithm for computing the Wasserstein barycenter of arbitrary distributions. Our algorithm can operate directly on continuous input distributions and is optimized for streaming data. Our method is even robust to nonstationary input distributions and produces a barycenter estimate that tracks the input measures over time. The algorithm is semi-discrete, needing to discretize only the barycenter estimate. To the best of our knowledge, we also provide the first bounds on the quality of the approximate barycenter as the discretization becomes finer. Finally, we demonstrate the practical effectiveness of our method, both in tracking moving distributions on a sphere, as well as in a large-scale Bayesian inference task.
barycenter of the input measures [1], and should be thought of as an aggregation of the input measures which preserves their geometry. This particular aggregation enjoys many nice properties: in the earlier Bayesian inference example, aggregating subset posterior distributions via their Wasserstein barycenter yields guarantees on the original inference task [47].
If the measures µ j are discrete, their barycenter can be computed relatively efficiently via either a sparse linear program [2], or regularized projection-based methods [16,7,51,17]. However, 1. these techniques scale poorly with the support of the measures, and quickly become impractical as the support becomes large. 2. When the input measures are continuous, to the best of our knowledge the only option is to discretize them via sampling, but the rate of convergence to the true (continuous) barycenter is not well-understood. These two confounding factors make it difficult to utilize barycenters in scenarios like parallel Bayesian inference where the measures are continuous and a fine approximation is needed. These are the primary issues we work to address in this paper.
Given sample access to J potentially continuous distributions µ j , we propose a communicationefficient, parallel algorithm to estimate their barycenter. Our method can be parallelized to J worker machines, and the messages sent between machines are merely single integers. We require a discrete approximation only of the barycenter itself, making our algorithm semi-discrete, and our algorithm scales well to fine approximations (e.g. n ⇡ 10 6 ). In contrast to previous work, we provide guarantees on the quality of the approximation as n increases. These rates apply to the general setting in which the µ j 's are defined on manifolds, with applications to directional statistics [46]. Our algorithm is based on stochastic gradient descent as in [22] and hence is robust to gradual changes in the distributions: as the µ j 's change over time, we maintain a moving estimate of their barycenter, a task which is not possible using current methods without solving a large linear program in each iteration.
We emphasize that we aggregate the input distributions into a summary, the barycenter, which is itself a distribution. Instead of performing any single domain-specific task such as clustering or estimating an expectation, we can simply compute the barycenter of the inputs and process it later any arbitrary way. This generality coupled with the efficiency and parallelism of our algorithm yields immediate applications in fields from large scale Bayesian inference to e.g. streaming sensor fusion.
Contributions. 1. We give a communication-efficient and fully parallel algorithm for computing the barycenter of a collection of distributions. Although our algorithm is semi-discrete, we stress that the input measures can be continuous, and even nonstationary. 2. We give bounds on the quality of the recovered barycenter as our discretization becomes finer. These are the first such bounds we are aware of, and they apply to measures on arbitrary compact and connected manifolds. 3. We demonstrate the practical effectiveness of our method, both in tracking moving distributions on a sphere, as well as in a real large-scale Bayesian inference task.
Title: Parallel Streaming Wasserstein Barycenters
Comments:
- This paper presents a new method for performing low-communication parallel inference via computing the Wasserstein barycenter of a set of distributions. Unlike previous work, this method aims to reduce certain approximations incurred by discretization. Theoretically, this paper gives results involving the rate of the convergence of the barycenter distance. Empirically, this paper shows results on a synthetic task involving a Von Mises distribution and on a logistic regression task.
- I feel that the clarity of writing in this paper is not great. It would be better to clearly (and near the beginning of the paper) give an intuition behind the methodology improvements that this paper aims to provide, relative to previous work on computing the Wasserstein barycenter. This paper quickly dives into the algorithm and theory details (“Background”, “Mathematical preliminaries”, “Deriving the optimization problem”), without giving a clear treatment of previous work, and how this methods of this paper differ from this work. It is therefore is hard to see where the material developed in previous work ends and the new methodology of this paper begins. A simple description (early on) of this method, how it differs from existing methods, and why it solves the problem inherent in these existing methods, would greatly increase the clarity of this paper. 
- Furthermore, it would be nice to include more motivation behind the main theoretical results that are proved in this paper. It is hard to get a grasp on the usefulness and contribution of these results without some discussion or reasoning on why one would like to prove these results (e.g. the benefits of proving this theory).
- Finally, I do not think that the empirical results in this paper are particularly thorough. The results in Table 1 are straightforward, but these seem to be the only empirical argument of the paper, and they are quite minimal.