Dimensionality Reduction for Wasserstein Barycenter
Authors: Zachary Izzo
Sandeep Silwal
Samson Zhou
Venue: NeurIPS 2021 Poster (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.

Dimensionality Reduction for Wasserstein Barycenter
Zachary Izzo Stanford University zle. wwzzzzzgaall.com
Sandeep Silwal
Samson Zhou Carnegie Mellon University sams samsonzhou@gmail. com
MIT silwal @mii t. edu
Abstract
The Wasserstein barycenter is 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 barycenter suffer an exponential dependence on the dimension d of the underlying space of the distributions. In order cope with this "curse of dimensionality"" we study dimensionality reduction techniques for the Wasserstein barycenter problem. When the barycenter is restricted support of size n, show that randomized dimensionality reduction can be used to map the problem to space of dimension (log n) independent of both d and k, and that 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.
an approximate
Introduction
The Wasserstein barycenter (WB) a popular method in statistics and machine learning for summa- rizing data from multiple sources while capturing their underlying geometry [A]] The problem defined as follows. Suppose we have a collection of data, represented as discrete probability distributions , 14 on Rd, Given set of non-negative weights λI,..., A that sum to 1, and a class P of probability distributions on Rd, a Wasserstein barycenter under the L, objective for parameter p > 0 is probability distribution v E P that minimizes
Σ W,Ws(,
where Wo(n,, v)) is the p-Wasserstein distance.
The Wasserstein barycenter is a natural quantity that captures the geometric notion of centrality among point clouds, as it utilizes the optimal transport distance [BT97] between a number of ob- served sets. Thus, Wasserstein barycenters have been extensively used in machine learning [SLD18], data sciences [RU02, EHJK20), image processing [RGT97], computer graphics [PW09], and statisties [Vil08], with applications in constrained clustering [CD14, HNY+ 17], Bayesian learn- ing (SLD18), texture mixing [RPDB1] I], and shape interpolation [SDGP+15]
Unfortunately, the problem is Na-hard to compute [AB21, BP21] and many algorithms that even approximate the Wasserstein barycenter suffer from large running times, especially the datasets are high dimensional [MC19]. Indeed, [ABA21] recently gave an algorithm that computes the
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Wasserstein barycenter using runtime that depends exponentially on the dimension, thus suffering the "curse of dimensionality."
To alleviate these computational constraints, we consider dimensionality reduction for computing the Wasserstein barycenter. Dimensionality reduction can be used to improve the performance of downstream algorithms on high dimensional datasets in many settings of interest, e.g., see the survey [CG15]. In the specific case of Wasserstein barycenters. dimensionality reduction has several practical and theoretical benefits, including lower storage space, faster running time in computing distances, and versatility: it can be used as a pre-proeessing tool and combined with any algorithm for computing the Wasserstein barycenter.
1.1 Our Results
In this paper, we study dimensionality reduction techniques for computing a Wasserstein barycenter of discrete probability distributions. Our main results show that it is possible to project the distributions into low dimensions while provably preserving the quality of the barycenter. A key result in dimensionality reduction is the classical Cehnsoon-Lnndnsttenn projecting a dataset of N points into roughly Olog N) dimensions is enough to preserve all pairwise distances.
(JL) lemma [JL84], which states that
Using the JI. lemma, we first show that we can assume the distributions lie Ollo(nn)) dimensions, where & is the number of input distributions whose barycenter we are computing, is the size of the support of the barycenter, and each of the l input distributions has support size poly(n). For p = 2. there exists closed form for the cost of any candidate barycenter in terms of the pairwise distances of the points in the input distributions. Thus it is straightforward to see that our bound results from the fact that there are k. poly(n) total points masses in the union of all the distributions and therefore, projecting them into a dimension of size Olog(( /k poly(n)) = (log((nk) suffices to preserve all of their pairwise distances. However for p # 2, closed form for the optimal cost longer exists, so preservation of all pairwise distances is insufficient. Instead, we make use of a Lipschitz extension theorem, namely the Kirszbraun theorem, which allows to "invert" the dimensionality reduction map and argue the preservation of the cost of the Wasserstein barycenter under a general byp objective. For more details, see Section 3.
Dimensionality reduction independent of k. While the JL lemma is known to be tight [LNI6, LN17], it is possible to improve its dimensionality guarantees for specific problems, such as various formulations of clustering [CEM* 15, BBC 19, MMR19] Indeed, our main result is that we can achieve dimension bound beyond the Ollog(nk) (nk)) bound that follows from the JL lemma and Kirszbraun theorem. We show that it suffices to project the support points onto O(log n) dimensions, which is independent of the number of distributions k. In fact, projecting the points supported by the distributions onto O(log n) dimension preserves the cost of the objective (1) for any distribution v supported at most n points (Theorem 4.1). The algorithmic application of this theorem is that one can take any approximation algorithm or heuristic for computing the Wasserstein barycenter and combine it with dimensionality reduction. simplification of our theorem is stated below where we omit some parameters for clarity.
show stronger statement that
Theorem 1.1 (Theorem 4.1 Simplified). Let ji...., 14 be discrete probability distributions on IRd such that (supp(ra)) S poly(n) for all i. There exists a dimensionality reduction map π .R + Rm for m = Olog n.) such that projection under N. preserves the cost of objective (f)for any u supported on at most n points.
The result is surprising because the projected dimension is independent of the number of input distributions k which could be significantly larger than n. Thus the random projection map æ can no longer even guarantee the preservation of significant fraction of pairwise distances between the support points of the distributions. Our main tool is a "robust" Lipschitz extension theorem introduced in [MMR19] for k-means clustering. We adapt this analysis to the geometry of the Wasserstein barycenter problem.
Optimality of dimensionality reduction. We complement our upper bound results by showing that our dimension bound of log dimensions is tight if a random Gaussian matrix is used as the projection map. We also show that the JL lemma is tight for the related problem of computing the optimal transport between two distributions with support of size n. More specifically, we give a lower bound showing that Q(log n) dimension is needed for a random projection to preserve the optimal

---

transport cost. Thus our results show a separation between the geometry of the optimal transport problem and the geometry of the Wasserstein barycenter problem, as we overcome the JL bound in the latter.
Hardness of approximation. In addition, we also show the NP-hardness of approximation for the Wasserstein barycenter problem. Namely, we show that it is NP-hard to find an approximate barycenter that induces cost that is within a factor of 1.0013 of the optimal barycenter we restrict the support size of the barycenter. This complements recent work of [AB21, BP21], who showed that computing sparse Wasserstein barycenters NP-hard.
Coresets for Wasserstein barycenters. An alternate way to reduce the complexity of datasets is through the use of coresets, which decrease the effective data size by reducing the number of input points rather than the input dimension d. If the number of input distributions & is significantly larger than the support size n, we show that there exists weighted subset C of roughly poly(n) distributions, so that computing the optimal barycenter on C' is equivalent to computing the optimal barycenter on the original input up small approximation loss. Hence, can potentially be much more efficient to use the subset C' in downstream algorithms involving Wasserstein barycenters. Moreover, the coreset is not mutually exclusive with our techniques for reducing the ambient dimension d. Our techniques show that we can simultaneously reduce both the size of the input distribution k: and the dimension of the data, while preserving the optimal clustering within small approximation factor.
In Supplementary Section E, we also show a connection between the Wasserstein barycenter problem and constrained low-rank problems. This class of problems includes examples such as the singular value decomposition (SVD) and k-means clustering. While this connection does not yield any improved results, it classifies the Wasserstein barycenter and this classification could have further applications in the future.
member of a general class of problems,
Experiments. Finally, we present experimental evaluation of our proposed methodology. Note that our results imply that we can use dimensionality reduction in conjunction with any Wasserstein barycenter algorithm and still roughly retain the approximation guarantees of the algorithm used. Specifically, we give examples of real high dimensional datasets such that solving the Wasserstein barycenter problem in reduced dimension leads to computational savings while preserving the quality of the solution. Our experiments in Section 7 demonstrate that on natural datasets, we can reduce the dimension by 1-2 orders of magnitude while increasing the solution cost by only 5%. We also empirically test our coreset construction. Our method both reduces error and requires fewer samples than simple uniform sampling.
1.2 Related Work
[AB21, BP21] showed that computing sparse Wasserstein barycenters is NP-hard; hence, most of the algorithmic techniques focus computing approximate Wasserstein barycenters that induce cost within an additive E of the optimal cost. [AC1 1b] first considered approximating Wasserstein barycenters when either (1) the distributions TP only have discrete support on R, (2) k 2, (3) the distributions μ4; are all multivariate Gaussians in Rd. Although there is a line of research that studies the computation of barycenters of continuous distributions, e.g. [ÁDCM16, CMRS20), we focus discrete input distributions. For discrete input distributions, the majority of the literature can be categorized by its assumptions of the support of the barycenter [ABA21].
Fixed-support. The "fixed-support approximation" class of algorithms assume that the support of the barycenter among fixed set of possible points. then remains for the algorithms to solve a polynomial-size linear program associated with the corresponding set [CD14, BCC+15, 15, COO15, SCSJ17, KTD* 19, LHC*20]. Unfortunately, the set of possible points must often be an E-net over the entire space, which results in a size proportional to 1/e" that suffers from the curse of dimensionality. Nevertheless for constant dimension, the algorithms typically have runtime (nly(n, k D/E), where D is an upper bound on the diameter of the supports of the input distributions. This is further improved by an algorithm of [ABA21] that achieves runtime poly(n, k. log((De)))
Free support. A separate class of algorithms do not make assumptions about the possible support of the optimal barycenter. These "free-support algorithms" instead optimize over the entire set of up candidate barycenters, which can be as large as n* in quantity. Thus these algorithms, e.g.. [CD14, LSPC19], either use exponential runtime or number heuristics that lack theoretical guarantees. [ABA21] showed how to explore the n possible points in polynomial time for fixed d.
3

---

2 Preliminaries
Notation. For positive integer n. denote [-]:= := (1,2,....,n) We use j1.--.: to denote the k distributions whose Wasserstein barycenter we wish to compute. While the Wasserstein barycenter problem is well defined for continuous distributions, in practice and in actual computations, the distributions are assumed to be discrete distributions that are supported on some number of point masses. This is also the assumption we make. More specifically, we assume that each of the distributions 14; are discrete distributions supported on at most T S n points where C' is a fixed constant. That is, 14 Σj= (( where 5, is a delta function at x and a(x: is the weight assigned to point z in its corresponding 14, We note that if there is some point in the support of more than one of the μas, then the weight function a may not be well-defined. Instead, we implicitly assume that a = a(x, i) is function of both the point and the distribution from which it comes, but we suppress this dependence on i for notational clarity.
The distribution v denotes a candidate for the Wasserstein barycenter of the μ4:- We write v = Σj=1 05 In general, an actual Wasserstein barycenter (in the sense of minimizing the objective (1) over all possible of any support size) may have support size up to U supp(n))] [ABM16]. Throughout this paper, we will restrict ourselves to computing (approximate) barycenters of support size most n. When we refer to optimal barycenter, we mean a distribution that minimizes the objective (1) within this restricted class.
Problem description. The goal is to compute a distribution E Rd, consisting of at most point masses, to minimize the objective (l). As previously mentioned, W, ( v) is the Wasserstein p-metric, defined as
( 1k= wll -
W W)
inf yET(un)
where I(u, 1/) is the set of all joint distributions with marginals μ and v (i.e. all couplings of μu and v) and || denotes the Fuclidean norm on Rd. When and v are discrete distributions, W, (u, w)? p-metric is the cost of the minimum cost flow from μ to V with edge costs being the Euclidean distance raised to the p-th power. For simplicity, we assume that the distributions (1, Al are weighted equally (each A 1/k in (I)) but our results hold in the general case as well. The most common choice of p is 2.
Description of v/ The barycenter can be characterized follows. Recall that is supported on the For the optimal coupling of each μ4 to v, let (z) denote the total weight sent
points v/ ,...,,", 1/" from x (in the support of one of the p.s) to v². (The same note about suppressing the dependence of wy on the distribution μi; from which x comes applies here.) Let S = { E U supp(μ;) wy (4r) >> 0} denote the set of all points in the j&s with some weight sent to wd Then given the set Sj and weighting function wy (-), we can reconstruct vá since it must minimize the objective
w,(x)|<x- IIP.
SE
Indeed if nd does not minimize this quantity, we can change and reduce the cost of (1).
Consider the case of p = 2. For a fixed j. (2) is just a weighted k-means problem whose solution is the weighted average of the points in Sy- To prove this, consider taking the gradient of (2) with respect to the k-th coordinate of w² Then setting it equal to 0 gives us that the k-th coordinate will be the weighted average of the k-th coordinates of the points Sy. That is, we have
va Exes, wz(x)x
Σ wy(x)x. SES)
(3)
Exes, w)(x) kb;
The second equality results from observing that in order for the wys to define a proper coupling, we have Σ/=1 w (x) (z)) for all x in the support of the 14s, and E ExEsupp(r.) w (x) by for all i, along with wy (x) 2 O. In particular, this implies that Exes, wy w)(x) = k6, for all = I,... n.
For arbitrary P. such concise description of أ³ is not possible. Therefore an alternate, but equivalent, way to characterize the distribution v is to just define the sets S, and weight functions w,(+) for 1 I j S n. This motivates the following definitions.

---

Definition 2.1. A solution (S.w)= = (Si...., Sp.............. wi....,. w)) is a valid partition as described previously (meaning that these partitions come from the optimal coupling between each pr fixed v), along with the corresponding weight functions wy(+).
X
Figure 1: Points of the same color belong to the same distribution. The sets Sj are denoted by the large black circles. Given the partitions Sj (denoted by large black circles) and associated weight functions Wj. can reconstruct the barycenter (denoted by crosses).
Definition 2.2. Let (S, w) be a solution. The cost of this solution, denoted cost is the value of the objective (1) when we reconstruct V from S and w and evaluate (1):
1 "
cost,(S) min + Σ w,(x)|<= v v'l.P.
j=IzEs;
Similarly for projection т, (T) denotes the value of the objective (1) when we first project each of the distributions T. using n, then compute using the original weights wj:
costp( = min ΣE (x))))l(=))-bb) j=IZES,
Note that each DJ E R We suppress the dependence of the cost w for notational convenience.
For the case p = 2, we can further massage the value of eá in (2). Let denote the weighted average of points in Sy (given by (3)). From our discussion above, we know that vt = F. After some standard algebraic manipulation, we can show that Exes, (r(x)]|æ nd 1|2 = Ezes, 2z)x)||=11 kb, ||
and Ex,ves, wy (x))) yll 2kb, Eves, -)(()|l=|l--bb) equations yields the following for the p = 2 objective.
Combining these
1
Σ J.VES,
()(z) - y/1?
v
Σ SES;
wy(x))<= 101
(4)
2kby
Dimension reduction. In this paper we are concerned with dimensionality reduction maps Rd - R.T that are JL projections, i.e., any dimensionality reduction map that satisfies the condition of the JL lemma. This includes random Gaussian and uub-Gaussian matrices [LN16, MMR19]. We are mainly concerned with making the projection dimension m as small as possible.
Consider any algorithm A that, given ,1, ... (4k. solves for some approximate or exact v minimizing the objective (I). We can combine any such A with dimensionality reduction by first projecting the point masses of the μ4 down to RT for some m < d and using A to compute some barycenter ō in IR Then, we can consider the solution (8, w) induced by (see Definitions 2.1 and 2.2) to reconstruct the appropriate in the original dimension Rd using the objective Eq. (2). Note that this objective is a convex program for any p 2 1 since we are given S, and w.).). For p 2 (which is the most common case), yå has a particularly simple form which is the weighted average of the points in Sj (see Eq. (3)). This procedure is outlined in Algorithm 1.
As a corollary of our results, algorithm takes time T'n, k d), then using dimensionality reduction as in the procedure outlined above takes time T(n, k. m) plus the time to perform the projection and reconstruct the barycenter using the solution S. The cost of running algorithm much more expensive than performing the projection, and the reconstruction step can also be solved efficiently since it is convex. In the case of p = 2, the reconstruction just amounts to computing n weighted means. Therefore for < d. we get significant savings since T(n, k. m) << T(n,k, d).
is usually
5

---

Algorithm Using dimensionality reduction with any algorithm A for computing WB
Require: & discrete distributions " , with point masses dimension F4, projection dimen- sion m, algorithm A
Project the point masses of each distribution μu to dimension In using a JL projection
2: Use algorithm A to solve (or approximately solve) the Wasserstein barycenter problem in I" to get distribution
//D is a discrete distribution 3: Let (S, w) be the solution that partitions the the point masses of the distributions as described in Definition 2.1
4: for each Sj E ES do
5: Solve for y minimizing Exes, w)(x))|x - wlp //This is convex program for
p 2 For p= p=2, 2, vf is just the weighted average of points in Sg.
6: end for
7: Output the distribution v supported on v², and where v has the same weight as io
3 Reduction to O(log(nk)) Dimensions
We first show that it suffices to project the point masses the input distribution into Olog((n)) dimensions and guarantee that the cost of any solution is preserved. Note that results hold simultaneously for all solutions. We first state the = 2 case.
Theorem 3.1. Consider JL projection ग from I& to R. for m O(log(kk/0)/²²) Then
(cost2(TT)) E [[1 - =)² cS) + (S) for all solutions S) 2 1 --6. 6.
Proof. The proof follows from the solution decomposition given in (4) if we condition on all the pairwise distances being preserved which happens with probability 6.
A decomposition similar to (4) does not exist for p * 2. To prove an analogous theorem for + 2, we need the following Lipschitz extension theorem which roughly allows us "invert" dimensionality reduction map.
Theorem 3.2 (Kirszbraun Theorem [Kir34].. For any D) C Rm" let f D - Rd be an L-lipschitz
function. Then there exists some extension f I= Rdd of f to the entirety of T such that [a) = =(xx) for all D and f is also T.-Lipschitz.
The Kirszbraun theorem allows us prove Theorem 3.1 for general p with dimension bound of m = (llog(nk/8)pp) /e²) (see Theorem A.1 in Supplementary Section A).
The overview for the proof strategy for the general * 2 case is as follows. First suppose that all the pairwise distances between the support points all the distributions are preserved under the projection map up to multiplicative error 1 s. This event happens with probability at least 1 - 8. We then consider the map ʃ IRm + Idd that maps each of the projected points to its original counterpart in Rd. Note that the map is from the smaller dimension to the larger dimension d. On the support points, we know that ff is (1 + E)-Lipschitz by our assumption above.
Now if the projection caused the cost of T to decrease significantly, then using the Kriszbraun theorem, one could "lift" the corresponding barycenter ī from the projected dimension to the original dimension using the extension map f. Then since f is Lipschitz, this lifted barycenter (i) plugged into Eq. (2) would subsequently have cost smaller than the original barycenter that corresponds S in the original dimension. This is a contradiction in light of Eq. (2) and the description of v given in Section 2. Note that the exact description of f does not matter for the analysis, just that such a map exists. complete, rigorous proof can be found in the supplementary section.
4 Optimal Dimensionality Reduction
We now present our main theorem which improves the guarantees of Theorems 3.1 and A.I.

---

Theorem 4.1. Let pi..... Pk be discrete probability distributions on IR such that (supp([s)) S poly(n) for all i. Let d 2 1, e, E (0, 1), and p 2 1. Let d.m. Rd - Ih be a family of random
JL. maps with m = 0 Ex log Eb
Then we have,
P (costp((5) E [[1 - 8) cost,(S), (1+e) cos) for all solutions S) 2
We now give an overview of the proof strategy for Theorem 4,1, deferring all technical details to the supplementary section. Ideally, one would like to use a strategy similar to the proof of Theorem A. The key bottleneck is that when we project down to the specified in Theorem 4. 1, a large number of pairwise distances between the support points of the l distributions can be distorted (since we are projecting to a dimension smaller than Ollog(nk) (nk))) Therefore, the Kirszbraun theorem cannot apply as the map f described in the proof strategy of Theorem A. is no longer Lipschitz on the support points.
To overcome this barrier, we generalize an approach of [MMR19], who achieved the optimal dimen- sionality bounds for k-means clustering beyond the naïve JL bound by defining distortion graph on the set of input points, which has an edge between each pair of points if their pairwise distance is distorted by at least a (1+ --fccto under the random projection map They show that the distortion graph everywhere sparse, i.e., each vertex has small expected degree in the distortion graph, which implies "robust" Kirszbraun theorem (for their particular problem of k-means clustering). Namely, R in the projected space such that a large fraction of the distances from the pre-image f (v) to the input points in 12d are
there exists an extension map Rd - IR" and specific point
preserved. Moreover, the input points whose distance to f (v) is not preserved can be shown to contribute small eπror to the k-means clustering cost.
The dimensionality reduction maps of Theorem 4.1 generally require multiplication by a dense matrix of (scaled) subgaussian variables. In the Supplementary Section, we show that "faster" dimensionality reduction maps can also be used by providing a trade off between the projection runtime and the dimension m. Note that in practice, performing the projection is extremely cheap since only need to perform one matrix multiplication, which is highly optimized. Therefore the cost of any algorithm for Wasserstein barycenter will typically outweigh the cost of computing the projection.
4.1 Dimensionality Reduction Lower Bounds
In this section, we state lower bounds on the projection dimension m for the Wasserstein barycenter problem. Theorem 4.2 shows that Theorem 4.1 is tight up to constant factors.
Theorem 4.2. Consider the setup of Theorem 4.1. Any Gaussian matrix used as a dimension reduction map that allows a (1 + -aapproxximation to the optimal Wasserstein barycenter requires dimension Vloog n.
We also prove that one cannot do better than the naïve JL bound for the related problem of computing the optimal transport between two discrete distributions with n point masses each. This is in contrast to the case of Wasserstein barycenter where we were able to overcome the bound that comes from the JL lemma alone. Theorem 4.3 shows that the optimal solution in the projected dimension can induce a poor quality solution the original dimension the projection dimension is smaller than log n.
Theorem 4.3. There exists point sets A, B C TRd with |A| = [ = n and matching cost M between o(log n) dimensions using an appropriately
them, such that if randomly projected down scaled Gaussian random matrix, the pull back cost of the optimal matching in Ihth is at least w w(M).
In addition, we prove related theorem which states that the cost of the optimal transport is heavily distorted if we project to fewer than log n dimensions.
Theorem 4.4. There exists point sets A, B C Rd with = = n and matching cost M between them, such that if randomly projected down to o(log n) dimensions using an appropriately scaled Gaussian random matrix, the cost of optimal matching in R.m o(M) with probability at least 2/3.
See Supplementary Section D for full proofs.
7

---

5 Coresets
In this section, we give a coreset construction for Wasserstein barycenters. Our goal is to reduce the number of distributions k to only depend polynomially
We first define notion of coresets.
Definition 5.1 (Coreset). Fix p 1. Let C and M be two sets of distributions in R where all distributions consist of poly(n) point masses. is called an e-corset the set of distributions M if there exist weights We for E C such that for all distributions v of support size at most it holds that
(1 -
Σ We W)
1 Σ M
W( S (1 +
Σ dEC
We W(c,p).
GEC
NEM
The main result of this section is the following theorem.
Theorem 5.2 (Theorem C.6 simplified). Let M be a set of discrete distributions in Rdd, each supported on at most (nly(n) point masses. There exists a weighted subset K C M of size poly(n, d)/c²a that satisfies Definition 5.1 for p = O())
To prove Theorem 5.2, we follow the "importance sampling" by sensitivities framework in conjunction with using structural properties of the Wasserstein barycenter problem itself. The sensitivity sampling framework has been successfully applied achieve corsets for many problems in machine learning (see the references in the survey [BLK17]). Note that we have not attempted to optimize the constants our proofs and instead focus on showing that & can be reduced poly(n, d) for simplicity. The formal proof of Theorem 5.2 is deferred to the supplementary section.
We now describe the high level overview of the proof. We form the set C by sampling distributions in M with replacement based on their "importance" or contribution to the total cost. The notion of importance is formally captured by the definition of sensitivity.
Definition 5.3 (Sensitivity). Consider the N of all possible barycenter distributions with support size at most n. The sensitivity of a distribution E M is defined as
WH,
(p)) sup VEN
Epee W(d,
[
The total sensitivity is defined as G= M EpE
To see why such notion is beneficial, consider the case that one distribution μ consists of point masses that are outliers among all of the point masses comprising the distributions in M. Then is clear that we must sample with a higher probability if we wish to satisfy the definition of coreset. In particular, we sample each distribution in M with probability proportional to (an upper bound on) its sensitivity. Using a standard result in coreset construction, we can bound the size of the coreset in terms of the total sensitivity and a measure of the "complexity" of the Wasserstein barycenter problem which is related to the VC dimension. In particular,
utilize the notion of psuedo dimensoo.
Definition 5.4 Deeeud-Dimmnnonn, Definition 9 [LFKF18]). Let X be ground set and F be set of functions from X to the interval [0, 1]. Fix a set S = {21,- ... Fn C X, a set of reals numbers R {ri,....,r .Inl with r, [0, and function f E F. The called the induced subset of 5 formed by ff and R. The set S with associated values is shattered by F if {S, IS [ FJ 2n. The pseudo-dimension of F the cardinality of the largest shatlered subset of (or 0o).
S {x, E S f(x.) 2 rs) is
The following theorem provides a formal connection between the size of coresets and the notion of sensitivity and suuedodiieensoo. Note that the statement of the theorem is more general and applics to a wider class of problems. However, we specialize the theorem statement to the case of Wasserstein Barycenters.
Theorem 5.5 (Coreset Size, Theorem 2.4.6 in [Lan18], Theorem 2.3 in [BLK17] for the case of Wasserstein Barycenters). Let E > 0 and (0, 1). Let s M - 22 denote any upper bound function on the sensitivity o(.) defined in Definition 5.3 and let S = [M Epeas s(u). Consider set K of [K] samples of M with replacement where each distribution p E M is sampled with probability q(μ)= s(p)/(|M) S) and each sampled point is assigned the weight 1/(MM -KI-((K)) Let F.
8

---

denote the set of functions
W
Sq(+) ΣpEM W vip LUEN
where N is the set of all possible barycenter distributions with support size at most Let d' denote the pseudo-ding dimension of F. Then the set K (along with the associated weights) satisfies Definition
5.1 with probability at least - 5 if
cS [K| 2 6²
( d' dolog S + log 1
where is some absolute constant.
Thus, the bulk of our work lies in bounding the sensitivities and psuedo-dimension. For the former quantity, we exploit the fact that the Wasserstein distance a metric. The latter requires us tools from statistical learning theory which relate the VC dimension of a function class to its algorithmic complexity (see Lemmas C.3 and C.4). Full details given in Supplementary section C.
Other Theoretical Results
We now present some additional theoretical results pertaining Wasserstein barycenters. Our first result is that Wasserstein barycenters can be formulated as a constrained low-rank approximation problem. This class of problems includes coputing the SVD and k-means clustering [CEM+ 15]. Formally, we prove the following theorem.
Theorem 6.1. Given discrete distributions pi.... E IR with support size at most n, consider the problem of computing the Wasserstein barycenter with support size most n for the p objective. There exists matrix A E gnkxd and set S of rank orthogonal projection matrices in Rakk such that the first problem is equivalent computing
p'
argmin || TAPPA|I 1| PES
The proof of Theorem 6.1 is given in Section E.
We also prove the following NP hardness result Section F which complements the hardness results in [AB21, BP21].
Theorem 6.2. It is NP-hard to approximate an optimal Wasserstein barycenter of fixed support size up to multiplicative factor 1.0013.
Experiments
In this section, we empirically verify that dimensionality reduction can provide large computational savings without significantly reducing accuracy. We use the following datasets in our experiments.
FACES dataset: This dataset is used in the influential ISOMAP paper and consists of 698 images of faces in dimension 4096 [TSL00]. We form A = 2 distributions by splitting the images facing to the "left" versus the ones facing "right." This results in 350 uniform point masses per distribution.
We subsample 10 images from the MNIST test dataset (dimension 784). We split the images by their digit class which results in k 10 distributions with 103 uniform point masses each 784
MNIST dataset:
Experimental setup. We project our datasets in dimensions d ranging from d = 2 to d 30 and compute the Wasserstein barycenter for p the barycenter to be at most 5 points in 4046 (since the barycenter should intuitively return an "interpolation" between the left and right facing faces, should not supported on too many points). For MNIST we limit the support size of the barycenter to be at most 40. We then take the barycenter found in the lower dimension and compare its cost the higher dimension (see Algorithm against the Wasserstein barycenter found in the higher dimension.
For FACES, we limit the support size of

---

We use the code and default settings from [Ye19] to compute the Wasserstein barycenter; this implementation has been applied in previous empirical papers [YWWL17]. While we fix this implementation, note that dimensionality reduction is extremely flexible and can work with any algorithm or implementation (see Algorithm 1) and we would expect it to produce similar results.
Results. Our results are displayed in Figure 2. We see that for both datasets, reducing the dimension 30 only increases the cost of the solution by 5%. This 1-2 orders of magnitude smaller than from the original dimensions of 784 and 4096 for MNIST and FACES respectively. The average time taken run the Wasserstein barycenter computation algorithm in d 30 was 73% and 9% of the time taken to run in the full dimensions respectively.
784
FACES, 4096
1.14
1.40
Batio of Solution Quality
Ratoo of Solution Quality 10
1.35
4.12
1.30
1.25
2.30
1.20
1.15
1.68
1:10
1.05
1:00
15
30
15
30
Projection Dimension
Projection Dimension
(a) MNIST
(b) Faces
Figure 2: Ratio of the quality of solution found in the lower dimension versus the original dimension. Result displays average of 20 independent trials and #1 standard deviation shaded.
Coreset experiments. Our coreset result reduces the number of distributions k through sensitivity (importance) sampling. We created synthetic dataset with large k but small and d emphasize the advantage of sensitivity sampling over uniform sampling. We have k = 50, 000 distributions that each consists of single point mass in E. The first k distributions are all supported at the origin while one distribution is supported at = k. We consider the p 2 case and limit the support size of the barycenter also be 1. Let costang (v) denote the cost of on the original objective (1) and let costome (V) the cost of (1) when evaluated on a coreset. We record the relative evaluated at òg, i.e. single unit point mass at a, for x 0, 1, 10. We then average the results across 10 trials each. As (the point on which the query distribution supported) grows bigger, the associated cost became bigger, hence decreasing the relative error. Other query locations displayed the same trend. See Figure 3 for more details.
error ccostcare (v) costorig costanggu))// [/)l///ostang(((
% error at query
Method
# of samples
100
10
2 0
Uniform sampling
1000
0.986
9.087
49.998
100
Sensitivity sampling
10
0.0040
0.0036
0.0020
O
Figure 3: Even with much fewer samples, sensitivity sampling outperforms uniform sampling for number of query locations, averaged across 10 repetitions.
Acknowledgments
Sandeep Silwal was supported part by NSF Graduate Research Fellowship Program. Samson Zhou was supported by a Simons Investigator Award of David P. Woodruff.
References
[AB99] Martin Anthony and Peter Bartlett. Neural Network Learning: Theoretical Founda- tions. Cambridge University Press, 1999.
[AB21] Jason M. Altschuler and Enric Boix-Adserà. Wasserstein barycenters are np-hard to compute. CoRR, abs/2101.01100,, 2021.
10

---

[ABA21] Jason M Altschuler and Enric Adix-a Wasserstein barycenters can be computed in polynomial time fixed dimension. Journal of Machine Learning Research, 22(44):1- 19, 2021.
[ABM16] Ethan Anderes, Steffen Borgwardt, and Jacob Miller. Discrete wasserstein barycenters: Optimal transport for discrete data. Mathematical Methods of Operations Research, 84, 10 2016.
[AC09] Nir Ailon and Bernard Chazelle. The fast johnson-lindenstraus transform and approxi- mate nearest neighbors. SIAM Comput., 3911)302-33 2009.
[AC1Ia] Martial Agueh and Guillaume Carlier. Barycenters in the wasserstein space. SIAM J. Math. Analysis, 43:904-924, 01 2011.
[AC1Ib] Martial Agueh and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904 924, 2011.
[ÁDCM16] Pedro C Alvarez-Esteban, Del Barrio, JA Cuesta-Albertos, and C Matrán. A fixed- point approach to barycenters in wasserstein space. Journal of Mathematical Analysis and Applications, 44 441(4):444-76) 2016.
[BBC* 19] Luca Becchetti, Mare Bury, Vincent Cohen-Addad, Fabrizio Grandoni, and Chris Schwicgelshohn. Oblivious dimension reduction for k-means: beyond subspaces and the Sebnsoon-indenstrrann lemma. In Proceedings of the 51 Annual ACM SIGACT Symposium on Theory of Computing, STOC, pages 1039-1050, 2019.
[BCC+15] Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré. Iterative bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2)A1111-AI 2015.
[BK117] Olivier Bachem, Mario Lucic, and Andreas Krause. Practical coreset constructions for machine learning, 2017.
[BP21] Steffen Borgwardt and Stephan Patterson. On the computational complexity of finding sparse wasserstein barycenter. J. Comb. Optim., 41(3);736-71, 2021.
[BT97] Dimitris Bertsimas and John Tsitsiklis. Introduction to linear optimization, volume Athena Scientific Belmont, MA, 1997.
[CC06] Miroslav Chlebík and Janka Chlebiková. Complexity of approximating bounded variants of optimization problems. Theor. Comput. Sci., 354(3):320-338 2006.
[CD14] Marco Cuturi and Arnaud Doucet. Fast computation of wasserstein barycenters. In International conference on machine learning, pages 685-693, 2014.
[CEM+15) 15] Michael B. Cohen, Sam Elder, Cameron Musco, Christopher Musco, and Madalina Persu. Dimensionality reduction for k-means clustering and low rank approximation. In Proceedings of the Forty-Seventh Annual ACM on Symposium on Theory of Computing, STOC, pages 163-172, 2015.
[CG15] John P Cunningham and Zoubin Ghahramani. Linear dimensionality reduction: Survey, insights, and generalizations. The Journal of Machine Learning Research, 16(1):2859- 2900, 20155
[CMRS20] Sinho Chewi, Tyler Maunu, Philippe Rigollet, and Austin J Stromme. Gradient descent algorithms for bures wassersten barycenters. In Conference on Learning Theory, pages 1304, 2020.
[C0015] Guillaume Carlier, Adam Oberman, and Edouard Oudet. Numerical methods for matching for teams and wasserstein barycenters. ESAIM: Mathematical Modelling and Numerical Analysis, 49(6):1211-162) 2015.
[EHJK20] Filip Elvander, Isabel Haasler, Andreas Jakobsson, and Johan Karlsson. Multi-marginal optimal transport using partial information with applications in robust localization and sensor fusion. Signal Processing, 171:107474, 2020.

---

[Fel20] Dan Feldman. Introduction to core-sets: an updated survey. CoRR, C.000444,001100938 2020.
[HNY** 17] Nhat Ho, XuanLong Nguyen, Mikhail Yurochkin, Hung Hai Bui, Viet Huynh, and Dinh Phung. Multilevel clustering via wasserstein means. In International Conference Machine Learning, pages 1501-1509. PMLR, 2017.
[IN07] Piotr Indyk and Assaf Naor. wearrst-niihhbor-preeeeiin embeddings. ACM Trans. Algorithms, 3(3):31-es, August 2007.
[JL84] William B Johnson and Joram Lindenstrauss. Extensions of lipschitz mappings into hilbert space. Contemporary mathematics, 26 (1962)-26)),, 1984.
[Kir34] M. Kirszbraun. Über die zusammenziehende und lipschitzsche transformationen. Fun- damenta Mathematicae, 22(1):77-108,, 1934.
[KMNII] Daniel M. Kane, Raghu Meka, and Jelani Nelson. Almost optimal explicit johnson- lindenstrauss families. In Approximation, Randomization, and Combinatorial Optimiza- tion. Algorithms and Techniques 14th International Workshop. APPROX, and 15th International Workshop. RANDOM. Proceedings, pages 228-639, 2011.
[KTD*19] 19] Alexey Kroshnin, Nazarii Tupitsa, Darina Dvinskikh, Pavel Dvurechensky, Alexan- der Gasnikov, and Cesar Uribe. On the complexity approximating wasserstein barycenters. In International conference machine learning, pages 3530-3540, 2019,
[Lan18] Harry Lang. Streaming Coresets for High Dimensional Geometry. PhD thesis, Johns Hopkins University, 7 2018
[LFKF18] Mario Lucic, Matthew Faulkner, Andreas Krause, and Dan Feldman. Training gaussian mixture models at scale via coresets. Journal of Machine Learning Research, 18(16)):1 25, 2018.
[LHC+20] Tianyi Lin, Nhat Ho, Xi Chen, Marco Cuturi, and Michael Jordan. Fixed-support wasserstein barycenters: Computational hardness and fast algorithm. Advances in Neural Information Processing Systems, 33, 2020.
[LN16] Kasper Green Larsen and Jelani Nelson. The beohsoo--lindenstrnne lemma is optimal for linear dimensionality reduction. In 43rd International Colloquiun on Automata, Languages, and Programming, ICALP, 2016.
[LN17] Kasper Green Larsen and Jelani Nelson. Optimality of the johnson-lindenstrane lemma. In 58th IEEE Annual Symposium on Foundations of Computer Science, FOCS, pages 633-638, 2017.
[LSPC19] Giulia Luise, Saverio Salzo, Massimiliano Pontil, and Carlo Ciliberto. Sinkhorn barycen- ters with free support via frank-wolfe algorithm. In Advances in Neural Information Processing Systems 32: Annual Conference Neural Information Processing Systems, pages 9-199, 2019.
[LSW17] Euiwoong Lee, Melanie Schmidt, and John Wright. Improved and simplified inapprox- imability for k-means. Inf. Process. Lett., 120:40 43, 2017.
[MC19] Boris Muzellec and Marco Cuturi. Subspace detours: Building transport plans that are optimal on subspace projections. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems, NeurIPS, pages 6914-6925, 2019.
[MR11] Konstantin Makarychev, Yury Makarychev, and Ilya P. Razenshteyn. Performance of johnson- linneons-indeesssruusss transform for k-means and k-medians clustering. Proceedings of the / Annual ACM SIGACT Symposium on Theory of Computing, STOC, pages 1037-1038, 2019.
[PW09] Ofir Pele and Michael Werman. Fast and robust earth mover's distances. In 2009 IEEE 12th international conference on computer vision, pages 460-467. IEEE, 2009.
12

---

[RGT97] Yossi Rubner, Leonidas Guibas, and Carlo Tomasi. The earth mover's distance, multi-dimensional scaling, and color-based image retrieval. In Proceedings of the ARPA image understanding workshop. volume 661, page 668, 1997.
[RPDB11] Julien Rabin, Gabriel Peyré, Julie Delon, and Mare Bernot. Wasserstein barycenter and its application to texture mixing. In International Conference on Scale Space and Variational Methods in Computer Vision, pages 435-446, 2011.
[RU02] Ludger Rüschendorf and Ludger Uckelmann. On the n-coupling problem. Journal of multivariate analysis, (12)2442-258, 2002.
[SCSJ17] Matthew Staib, Sebastian Claici, Justin Solomon, and Stefanic Jegelka. Parallel stream- ing wasserstein barycenters. pages 2647-2658, 2017.
[SDGP*15 15] Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphics (TOG), 34(4):1-11,
[SLD18] Sanvesh Srivastava, Cheng Li, and David B Dunson. Scalable bayes via barycenter in wasserstein space. The Journal of Machine Learning Research, 19(1):312-345 2018.
[TSL00) Joshua B. Tenenbaum, Vin de Silva, and John C. Langford. A global geometric framework for nonlinear dimensionality reduction. Science, 23365500) 2000.
[Vi108] Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008
[Wai19] Martin J Wainwright. High-dimensional statistics: non-asymptotic viewpoint, vol- ume 48. Cambridge University Press, 2019.
[Ye19] Jianbo Ye. Wbc-matlab. attps://githes wwm.bobyye/WBC_aatl..... 2019.
[YWWL17] Jianbo Ye, P. Wu, J. Z. Wang, and Jia Li. Fast discrete distribution clustering using wasserstein barycenter with sparse support. IEEE Transactions on Signal Processing, 65:2317-2332, 2017.
13


Reviews:
[{"review_id": "QsKlrWZfZhN", "reviewer": "Reviewer_W1WL", "summary": "Paper studies dimensionality reduction and coresets for the Wasserstein barycenter problem. The objects are distributions supported on n points in R^d, and the barycenter is a single distribution which minimizes sum of Wasserstein distances (or related quantities). The show that using random projections, they can reduce the dimension of the ambient space from R^d to R^{O(log(n))} while preserving the cost of optimal barycenter. They also give coresets from which one can compute the Wasserstein barycenter.", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "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. \n\nA few remarks: \n1. This problem can be interpreted as 1-mean (or 1-median) clustering, where distances are measured according to the Wasserstein distance.\n2. 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).\n\nThe 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.\n\nI 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. \n\nOne 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 (?).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "JF0TG0Wrmng", "reviewer": "Reviewer_e4YX", "summary": "This paper suggests a coreset construction scheme and a dimensionality reduction for the Wasserstein Barycenter, which can be thought of the $\\ell_p$-mean of $k$ discrete distributions where the mean here is a probability distribution.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "In this paper, a dimensionality reduction via JL has been proposed followed by a coreset construction, yielding a small weighted coreset. \nThe 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. \nUsually in the literature of coresets, when the weak triangle inequality concerning the loss functions holds, then one can attain a coreset easily. \n\nRegardless, I find the paper interesting. \n\nFurthermore, The authors suggested the first upper bound on the vc dimension concerning the problem.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "- The coreset size depends heavily on the size of the support of the given $k$ discrete probability distribution, i.e., $O\\left(n^8\\right)$. This bound actually comes from the bound on the VC dimension of the problem at hand. \nThe coreset becomes impractical when $n$ is large.  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 depencedy on the total sensitivity from near-linear to quadratic at the sample size (concerning the coreset)?\n\n- Experimental results:\n  a. 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) --> at the end this is what is important.\n\n  b.  why there are no graphs in the coreset comparison to uniform sampling (the reader wants to see phenomena, not just numbers) +  add variance to these graphs --> it is an important measure.\n\nNote-> it is almost clear that the importance sampling is better, however, I think there is a lot of experiments that should be done more to emphasize that.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}, {"review_id": "1uwmqT7JJr", "reviewer": "Reviewer_gFL6", "summary": "This paper justifies compute Wasserstein barycenter with dimension reduction techniques, which can lead to significant cost saving. Specifically, with the help of JL lemma, the original Wasserstein problem can be solved in a space of dimension $\\mathcal{O}(\\log n)$, where n is the specified barycenter support size. The reduced dimension is **tight**, and is not dependent on the original problem dimension and the number of input distributions. The authors also discuss the hardness of the barycenter approximation, the existence of a coreset with size poly($n$), and the connection between the Wasserstein barycenter problem and constrained low-rank problems. The authors empirically validate their claims using real datasets on dimension reduction and a synthetic dataset on coreset construction.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper flows well. The related concepts are well introduced, and the proofs are properly outlined. \n\nThe 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. \n\nThe proof techniques might shed lights on other related future works on dimension reduction.\n\n\nMy 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? \n\nIn 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?**\n\n\nSome minor comments:\n1.\tI 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.\n2.\tFigure 2 would be more straightforward if there is a horizontal line at ratio=1.\n3.\tFigure 3 should be a table, and the caption should locate on the top of the table.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "This paper mainly focuses on the mathematical and algorithmic properties of the Wasserstein barycenter problem, not do not focus on any particular downstream applications that may have negative social impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: 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)