Title: Dimensionality Reduction for Wasserstein Barycenter
Authors: Zachary Izzo
Sandeep Silwal
Samson Zhou
Venue: NeurIPS 2021
Keywords: wasserstein barycenter
dimensionality reduction

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.

TL;DR: We study randomized dimensionality reduction for the Wasserstein barycenter problem.

Review discussion:
This paper flows well. The related concepts are well introduced, and the proofs are properly outlined.
The idea of using dimension reduction to speed up the computation of Wasserstein barycenter is promising. This paper provides a comprehensive discussion of the applicability of dimension reduction – when the cost can be preserved and when cannot. The speedup from algorithm 1 should be significant.
The proof techniques might shed lights on other related future works on dimension reduction.
My questions mainly focus on the computation time. In the MNIST experiment, the dimension is reduced from 784 to 30, which is remarkable.
However, the average time taken is **73%** of without dimension reduction. I wonder why the time reduction is not very significant. What is the decomposition of the time usage in this particular experiment, i.e., what is the time cost of each step of algorithm 1? As claimed in line 207, the cost of running algorithm A is usually much more expensive than performing the projection – does this still hold for this experiment?
In contrast, the time usage for FACES is only 9% of the original one. What causes the difference? Larger original dimensions? Smaller support sizes of the input distributions and the barycenter? More generally, in what cases the speedup from dimension reduction is more significant?
Also, have you tested on different barycenter support sizes? **Is the current time reduction optimal with respect to the barycenter support size?**
Some minor comments:
1.	I guess “ratio of the quality of solution” in Figure 2 is the ratio between the costs (as implied by line 353)? It would be better if this is more clear in the caption.
2.	Figure 2 would be more straightforward if there is a horizontal line at ratio=1.
3.	Figure 3 should be a table, and the caption should locate on the top of the table.
No
Thank you for the positive comments on our paper. We have addressed the minor concerns that you noted, for the next version of the paper. We address the main questions as follows.
> "I wonder why the time reduction is not very significant."
We believe this is due to the decrease in sparsity after performing the random projection. In the original dimension, the input points are quite sparse whereas after a random projection, the input points become dense even if the dimension is reduced significantly.
> "As claimed in line 207, the cost of running algorithm A is usually much more expensive than performing the projection – does this still hold for this experiment?"
Yes.
In practice, a random projection can be done using a single matrix multiplication or even faster using FFT methods, such as in the `fast JL transform.’ Since matrix multiplication is a highly optimized operation and is only done once, this step is seemingly ‘for free’ and is many orders of magnitude faster than running a barycenter algorithm, even in the reduced dimension.
> "Also, have you tested on different barycenter support sizes? Is the current time reduction optimal with respect to the barycenter support size?"
For our experiments we only tested on the support size stated in the experimental section, though we have run multiple iterations to track the average performance (as well as central tendencies such as the variance). Since we sought an empirical proof-of-concept to support our main theoretical statements, we also did not attempt to optimize the running time of our algorithm. We believe using alternate support sizes result in similar performance gains.
> "Minor comments, I guess “ratio of the quality of solution” in Figure 2 is the ratio between the costs (as implied by line 353)? It would be better if this is more clear in the caption."
Thank you for the valuable feedback. Yes, the ratio is the ratio of costs of the solution found in the projected dimension (but evaluated in the original dimension) and the solution found in the original dimension. We have updated these comments for the next version of the paper.
Still, I am curious what is the decomposition of the time usage in this particular experiment, i.e., what is the time cost of each step of algorithm 1?
Using the Faces dataset and the experimental setup as described in Section 6, we get approximately the following times (averaged over some new runs performed today):
Time taken to perform projection to d = 30: 0.02 seconds
Time taken to run the WB solver on the projected datasets: 0.43 seconds
Time taken to reconstruct the solution (as per the loop in Line 4 in Algorithm 1): 0.17 seconds.
Contrast that with the time taken to run the WB solver on the original dataset: 7.81 seconds.
In this paper, a dimensionality reduction via JL has been proposed followed by a coreset construction, yielding a small weighted coreset.
The ideas behind applying these methods rely on the Lipschitz property of the cost function, whereas the coreset relied on the sensitivity sampling framework which in this paper, exploited the weak triangle inequality concerning the cost function.
Usually in the literature of coresets, when the weak triangle inequality concerning the loss functions holds, then one can attain a coreset easily.
Regardless, I find the paper interesting.
Furthermore, The authors suggested the first upper bound on the vc dimension concerning the problem.
No
Thank you for your comments, we are glad that you found the paper interesting.
> "Is your bound on the VC dimension tight? Would it be better to use the doubling dimension instead of the VC dimension, while increasing the dependency on the total sensitivity from near-linear to quadratic at the sample size (concerning the coreset)?"
We do not believe that the bound on the VC dimension is tight as our goal was to get the first coreset size independent of the number of distributions.
Therefore, we did not strictly optimize for the parameters of the coreset size.
In practice, it is not important to know the exact coreset size bound as we can perform importance sampling based on the available space and time resources. The doubling dimension parameterization of the sensitivity is an interesting future direction.
> "Experimental results"
Thank you for the valuable feedback. We have added graphs for the coreset experiments to the next version of the paper. We hope that if the presentation of the empirical results is your main concern, you would consider raising your score.
Can you please elaborate more, what are the graphs that you have added (x-axis, y-axis)?
Our graphs are standard plots for coreset papers: the x-axis represents the number of sampled distributions whereas the y-axis represents the relative error. The curve for the uniform sampling scheme is higher than that of importance sampling, signifying that importance sampling gives us smaller error using the same "budget."
Thank you very much for your reply, it would be very helpful if you can do the following experiment (I suggested previously):
Add graphs where the x-axis is the time took to compute a coreset + to solve the problem on the coreset (compared to uniform sampling) --> sometimes what is interesting is not the "budget" in terms of memory but rather the runtime to obtain an approximation.
Thanks for your inquiry. We have run experiments on our synthetic dataset that generates 50,000 points, to evaluate the relationships you suggested.
For sample size 10, the total time across 100 repetitions of uniform sampling was roughly $0.004$ seconds with variance $1.59\times10^{-7}$, while importance sampling was roughly $0.07$ seconds with variance $2.36\times10^{-6}$. The average error of uniform sampling was roughly $50\%$ while the average error of the coreset was roughly $0.02\%$.
For sample size 100, the total time across 100 repetitions of uniform sampling was roughly $0.02$ seconds with variance $3.31\times10^{-7}$, while importance sampling was roughly $0.09$ seconds with variance $1.08\times10^{-6}$. The average error of uniform sampling was roughly $50\%$ while the average error of the coreset was roughly $0.002\%$.
For sample size 1000, the total time across 100 repetitions of uniform sampling was roughly $0.21$ seconds with variance $4.34\times10^{-6}$, while importance sampling was roughly $0.26$ seconds with variance $3.63\times10^{-6}$. The average error of uniform sampling was roughly $25\%$ while the average error of the coreset was roughly $0.002\%$.
We will perform more extensive evaluations, incorporate the results into a table and add them into the next version of the paper.
The paper considers the Wasserstein barycenter problem: we receive k distributions supported on n points in R^d, and the task is to find a single distribution which is a good 'center' for these k distributions.
In particular, the center minimizes the sum-of-p-th powers of Wasserstein-p distances from the single distribution to the k other distributions.
A few remarks:
1. This problem can be interpreted as 1-mean (or 1-median) clustering, where distances are measured according to the Wasserstein distance.
2. The paper actually considers a constrained version of the problem, where the center distribution is also supported on at most n points in R^d. With this constrain, the problem can be interpreted as a n-cluster problem in R^d with a few additional constraints on the allowed clustering. At a high level, consider a barycenter which is specified by n points c_1, ..., c_n in R^d and n weights a_1, ..., a_n which specify the probability that we sample c_i.
Then, for each of the k distributions, think of them as uniform over a *multi-set* in R^d. The Wasserstein distance from a distribution x to the barycenter is given by an assignment of the points to the n centers with the additional constrain that for each center c_i, the total weight is a_i.
In other words, this is a n-clustering problem where we consider partitions of the total support with k constraints (one for each distribution).
The paper shows that one can reduce the dimensionality to O(log(n)/eps^2) dimensions with JL and preserve the cost up to 1+eps, and how to construct coresets with poly(n) distributions.
Furthermore, they show that O(log(n)/eps^2) dimensions are necessary for preserving the Wasserstein distance between two points (and hence this is a limit).
Finally, they show it is NP-hard to find an approximate Wasserstein barycenter supported on n points up to some constant factor.
I am confused about the following. The paper builds heavily on [Makarychev-Makarychev-Razenshteyn], who showed that in k-clustering problems, the dimension may be reduced to O(log(k)) and preserve the cost to *all* partitions of the points into k parts.
Furthermore, coreset results for k-clustering objectives usually have the property of being "strong coresets" in that they preserve cost to *all* partitions into k parts. If I understand the problem correctly, [MMR] directly implies Theorem 4.1? If I am misunderstanding some point, this should definitely be addressed in the introduction, since I don't see the challenges in utilizing [MMR] directly.
One additional comment: I don't see why the assumption that the distributions are supported on poly(n) points is necessary. It seems that the important aspect is that the actual barycenter is restricted to support-size n, a point which should be stressed in Theorem 1.1. It seems like an open problem to determine whether this assumption is actually needed (?).
No
Thank you for your comments. We hope that we can clarify a few points.
> "Random projections preserving the cost of optimal barycenter."
We note that our analysis shows that JL projections preserve the cost of any candidate barycenter, not only the optimal barycenter. (This is similar to the strong result given in [MMR].)
> "This problem can be interpreted as 1-mean (or 1-median) clustering, where distances are measured according to the Wasserstein distance."
We respectfully disagree. The Wasserstein barycenter problem and clustering are fundamentally different. This is because in k-means/median clustering, a center is allowed to serve arbitrarily many points from the input whereas in the barycenter case, we must have a weighted matching between the barycenter distribution and each of the input distributions. The matchings must respect the underlying probability distributions and thus a point mass that's part of an input distribution $\mu_i$ will not necessarily be matched to the closest point mass in the barycenter distribution.
> "If I understand the problem correctly, [MMR] directly implies Theorem 4.1?"
While we do utilize the tools introduced in MMR,  their results do not transfer over to the barycenter setting.
Therefore, we must overcome challenges such as (i) having weighted point masses, (ii) the fact that a single point mass of any input distribution can possibly be assigned to multiple point masses in the barycenter, (iii) respecting the matching constraints, and (iv), working under the Wasserstein (matching) metric. Please see Section 4, lines 261-278 where this is elaborated in significant details.
> "I don't see why the assumption that the distributions are supported on poly(n) points is necessary… It seems like an open problem to determine whether this assumption is actually needed (?)"
Our assumption can be equivalently understood as optimizing for the barycenter distribution that is supported on at most $\text{poly}(n)$ point masses if each of the input distributions are supported on at most poly(n) point masses. This is a natural assumption to make since it is common in practice to only search over barycenters with support size comparable (or smaller than) to the input distributions. In general, it is known that if we optimized over barycenters $\nu$ with arbitrarily large support sizes, we can assume that it is supported on at most $\| \sum_{i=1}^k \text{supp}(\mu_i) \|$ point masses as stated in Section 2.
Finally, we point out that Theorem 4.2 shows a lower bound that any dimensionality reduction map using a Gaussian matrix requires projected dimension at least $\Omega(\log n)$.
Thus if the distributions are supported on $N$ points, the projected dimension must be at least $\Omega(\log N)$, even if $N$ is larger than $\text{poly}(n)$.
We hope that these comments answer your questions and clarify the novel contributions of our paper, and if so that you would consider raising your score.
Sorry that I may not have communicated my concerns effectively. I am worried that the Wasserstein barycenter problem, where the barycenter has support at most $n$, can be massaged into a clustering problem with $n$ clusters. I'll try to expand on that; I am happy to change my score if I am mistaken, so please let me know if this is the case.
Let $\mu_1, \dots, \mu_k$ be the $k$ distributions in $\mathbb{R}^d$. Now, think of each $\mu_i$ as a *multi-set* $U_i$ of $\mathbb{R}^d$ of size $M$ (for some large enough $M$);
then the distribution $\mu_i$ simply samples a uniformly random element of $U_i$. By making $M$ large enough, we may view all $k$ distributions $\mu_1, \dots, \mu_k$ as multi-sets of size $M$.
In the Wasserstein barycenter problem, we want to find a distribution $\delta$, supported on points $c_1, \dots, c_n \in \mathbb{R}^d$ with weights $a_1, \dots, a_n \in [0,1]$ with $\sum_j a_j = 1$ which minimizes the sum of Wasserstein distances to the $k$ distributions $\mu_1, \dots, \mu_k$. If $\delta$ is fixed with points $c_1, \dots, c_n \in \mathbb{R}^d$ and weights $a_1, \dots, a_n \in [0,1]$ which sum to $1$, the Wasserstein distance between a distribution $\mu_i$ and $\delta$ is given by minimizing, over all assignments $f_i \colon U_i \to [n]$ where $(1/M) \sum_{u \in U_i} \mathbf{1}[f_i(u) = j] = a_j$ (this enforces the 'coupling' condition), the quantity
$\frac{1}{M} \sum_{u \in U_i} \| u - c_{f_i(u)}\|_2$.
Therefore, if we write out the optimization problem for the Wasserstein barycenter, we want to minimize
$\sum_{i=1}^k \frac{1}{M} \sum_{u \in U_i} \| u - c_{f_i(u)}\|_2$
subject to the following constraints:
(1) the points $c_1, \dots, c_n \in \mathbb{R}^d$.
(2) we have weights $a_1, \dots, a_n \in [0,1]$ with $\sum_{j} a_j = 1$.
(3) for each $i \in [k]$, we have a function $f_i \colon U_i \to [n]$, and for each $j \in [n]$, we have $(1/M) \sum_{u \in U_i} \mathbf{1}[f_i(u) = j] = a_j$.
Now, consider the multi-set $U$ given by the disjoint union of all $U_i$, for $i \in [k]$, and $f \colon U \to [n]$ be the function combining all $f_i$ (in particular, $f(u) = f_i(u)$ if $u$ is the element in $U_i$).
Then, the above problem can be re-written as minimizing
(*)     $(1/M) \sum_{u \in U} \| u - c_{f(u)}\|_2$
subject to the same constraints (1), (2), and (3). The $n$-median problem on the multi-set $U$ is the problem of minimizing the exact same expression (*) with the constraint (1) only.
Therefore, adding the constraints (2) and (3) changes the $n$-median problem to the Wasserstein barycenter problem, where barycenter is supported on $n$ points.
Furthermore, the constraints (2) and (3) limit the possible clusterings that we may consider;
however, the results in [Makarychev-Makarychev-Razenshteyn] hold for all possible clusterings of the multi-set $U$; which in particular, implies they hold for those clusterings with constraints (2) and (3).
For $k$-means clustering, the cost is generally defined so that the contribution of each point to the loss function is with respect to its closest center. For the Wasserstein barycenter problem, the cost induced by a point may not be with respect to the closest barycenter. As you noted however, the structure of the barycenter problem can be viewed as a number of distributions with uniformly weighted support on a multi-set subset of size N.
Indeed, we take this view both in Appendix E to establish a connection with constrained low-rank approximation, and in the proof of Theorem 4.1 to show the approximation guarantees by random JL maps. Unfortunately the discretization into subsets of size N does affect the assumptions required for the constrained LRA connection, but the discretization does not affect the assumptions required for the dimensionality reduction maps.
We would also like to highlight additional theoretical results that are independent of the work of [MMR19].
In particular, we derive coreset bounds which are independent of the number of input distributions. This can be viewed as another form of “dimensionality reduction”, as we introduce provable guarantees for an alternate method for reducing input size. We also give several lower bounds for Wasserstein Barycenter and related problems.
For example, we show it is NP hard to find better than a constant factor approximation in our setting. We also present dimensionality lower bounds for the closely related problem of optimal transport, i.e., matching between two distributions/point sets. Surprisingly, we show that it is not possible to do better than the JL bound for approximately preserving the optimal transport map, thus showing a separation between the geometry of the optimal transport problem and that of the Wasserstein barycenter problem.
Thanks for the reply. The idea to use dimension reduction for Wasserstein barycenter and the additional theoretical results, including hardness of approximation, coresets, and the lower bound are interesting.
However, the way that the paper is written, the main result (Theorem 1.1) and most of the technical work is a direct corollary of [MMR]. I think the writeup needs to be updated to reflect this before it is published.
Thank you for your constructive feedback. We will revise the manuscript to incorporate the proof that shows the elegant connection between [MMR] and one of our theorems and also highlight our other results.
We thank the reviewers for their valuable feedback. Answers are given in a response to each review.

Meta-review:
role: PC
data: comment: As pointed out by the reviewer W1WL, and discussed with the authors, the main result follows relatively directly from [MMR]. While the paper has other merits, the paper needs a major rewrite to address that.

Some reviewers were also not fully satisfied by the authors answer to the question on the time usage decomposition of the MNIST experiment.

Decision: Accept (Poster)