Dimensionality Reduction for Wasserstein Barycenter
Year: 2021

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 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.

heading: 1 Introduction
text: The Wasserstein barycenter (WB) is a popular method in statistics and machine learning for summarizing data from multiple sources while capturing their underlying geometry [AC11a]. The problem is defined as follows. Suppose we have a collection of data, represented as k discrete probability distributions µ1, . . . , µk on Rd. Given a set of non-negative weights 1, . . . , k that sum to 1, and a class P of probability distributions on Rd, a Wasserstein barycenter under the Lp objective for a parameter p > 0 is a probability distribution ⌫ 2 P that minimizes
kX
i=1
iWp(µi, ⌫) p, (1)
where Wp(µi, ⌫) 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 observed sets. Thus, Wasserstein barycenters have been extensively used in machine learning [SLD18], data sciences [RU02, EHJK20], image processing [RGT97], computer graphics [PW09], and statistics [Vil08], with applications in constrained clustering [CD14, HNY+17], Bayesian learning [SLD18], texture mixing [RPDB11], and shape interpolation [SDGP+15].
Unfortunately, the problem is NP-hard to compute [AB21, BP21] and many algorithms that even approximate the Wasserstein barycenter suffer from large running times, especially if the datasets are high dimensional [MC19]. Indeed, [ABA21] recently gave an algorithm that computes the
35th Conference on Neural Information Processing Systems (NeurIPS 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-processing tool and combined with any algorithm for computing the Wasserstein barycenter.

heading: 1.1 Our Results
text: 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 Johnson-Lindenstrauss (JL) lemma [JL84], which states that projecting a dataset of N points into roughly O(logN) dimensions is enough to preserve all pairwise distances.
Using the JL lemma, we first show that we can assume the distributions lie in O(log(nk)) dimensions, where k is the number of input distributions whose barycenter we are computing, n is the size of the support of the barycenter, and each of the k input distributions has support size poly(n). For p = 2, there exists a 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 O(log(k poly(n))) = O(log(nk)) suffices to preserve all of their pairwise distances. However for p 6= 2, a closed form for the optimal cost no 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 us to “invert” the dimensionality reduction map and argue the preservation of the cost of the Wasserstein barycenter under a general Lp objective. For more details, see Section 3.
Dimensionality reduction independent of k. While the JL lemma is known to be tight [LN16, 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 a dimension bound beyond the O(log(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, we show a stronger statement that projecting the points supported by the distributions onto O(log n) dimension preserves the cost of the objective (1) for any distribution ⌫ supported on 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. A simplification of our theorem is stated below where we omit some parameters for clarity. Theorem 1.1 (Theorem 4.1 Simplified). Let µ1, . . . , µk be discrete probability distributions on Rd such that |supp(µi)|  poly(n) for all i. There exists a dimensionality reduction map ⇡ : Rd ! Rm for m = O(log n) such that projection under ⇡ preserves the cost of objective (1) for any ⌫ 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 a significant fraction of pairwise distances between the support points of the k 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 n 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 ⌦(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 a cost that is within a factor of 1.0013 of the optimal barycenter if we restrict the support size of the barycenter. This complements recent work of [AB21, BP21], who showed that computing sparse Wasserstein barycenters is 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 k is significantly larger than the support size n, we show that there exists a 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 to a small approximation loss. Hence, it 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 d of the data, while preserving the optimal clustering within a 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 as a member of a general class of problems, and this classification could have further applications in the future.
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 a 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.

heading: 1.2 Related Work
text: [AB21, BP21] showed that computing sparse Wasserstein barycenters is NP-hard; hence, most of the algorithmic techniques focus on computing approximate Wasserstein barycenters that induce a cost within an additive " of the optimal cost. [AC11b] first considered approximating Wasserstein barycenters when either (1) the distributions P only have discrete support on R, (2) k = 2, or (3) the distributions µi 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 on 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 is among a fixed set of possible points. It then remains for the algorithms to solve a polynomial-size linear program associated with the corresponding set [CD14, BCC+15, COO15, SCSJ17, KTD+19, LHC+20]. Unfortunately, the set of possible points must often be an "-net over the entire space, which results in a size proportional to 1/"d that suffers from the curse of dimensionality. Nevertheless for constant dimension, the algorithms typically have runtime poly(n, k,D/"), 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(D/")).
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 nk in quantity. Thus these algorithms, e.g., [CD14, LSPC19], either use exponential runtime or a number heuristics that lack theoretical guarantees. [ABA21] showed how to explore the nk possible points in polynomial time for fixed d.

heading: 2 Preliminaries
text: Notation. For a positive integer n, we denote [n] := {1, 2, . . . , n}. We use µ1, . . . , µk 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 µi 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 µi are discrete distributions supported on at most T  nC points where C is a fixed constant. That is, µi = PT j=1 a(xij) xij , where x is a delta function at x and a(x) is the weight assigned to a point x in its corresponding µi. We note that if there is some point x in the support of more than one of the µis, then the weight function a may not be well-defined. Instead, we implicitly assume that a = a(x, i) is a function of both the point and the distribution from which it comes, but we suppress this dependence on i for notational clarity.
The distribution ⌫ denotes a candidate for the Wasserstein barycenter of the µi. We write ⌫ =Pn j=1 bj ⌫j . 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 | Sk
i=1 supp(µi)| [ABM16]. Throughout this paper, we will restrict ourselves to computing (approximate) barycenters of support size at most n. When we refer to an optimal barycenter, we mean a distribution that minimizes the objective (1) within this restricted class.
Problem description. The goal is to compute a distribution ⌫ 2 Rd, consisting of at most n point masses, to minimize the objective (1). As previously mentioned, Wp(µi, ⌫) is the Wasserstein p-metric, defined as
Wp(µ, ⌫) = inf 2 (µ,⌫)
✓Z
Rd⇥Rd kx ykpd (x, y)
◆1/p
where (µ, ⌫) is the set of all joint distributions with marginals µ and ⌫ (i.e. all couplings of µ and ⌫) and k · k denotes the Euclidean norm on Rd. When µ and ⌫ are discrete distributions, Wp(µ, ⌫)p p-metric is the cost of the minimum cost flow from µ to ⌫ with edge costs being the Euclidean distance raised to the p-th power. For simplicity, we assume that the distributions µ1, · · · , µk are weighted equally (each i = 1/k in (1)) but our results hold in the general case as well. The most common choice of p is p = 2.
Description of ⌫. The barycenter ⌫ can be characterized as follows. Recall that ⌫ is supported on the points ⌫1, . . . , ⌫n. For the optimal coupling of each µi to ⌫, let wj(x) denote the total weight sent from x (in the support of one of the µis) to ⌫j . (The same note about suppressing the dependence of wj on the distribution µi from which x comes applies here.) Let Sj = {x 2 Sk i=1 supp(µi) : wj(x) > 0} denote the set of all points in the µis with some weight sent to ⌫j . Then given the set Sj and weighting function wj(·), we can reconstruct ⌫j since it must minimize the objective
X
x2Sj
wj(x)kx ⌫ j k p. (2)
Indeed if ⌫j does not minimize this quantity, we can change it 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 Sj . To prove this, consider taking the gradient of (2) with respect to the k-th coordinate of ⌫j . 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 Sj . That is, we have
⌫j = P x2Sj wj(x)xP x2Sj wj(x) = 1 kbj X
x2Sj
wj(x)x. (3)
The second equality results from observing that in order for the wjs to define a proper coupling, we have Pn j=1 wj(x) = a(x) for all x in the support of the µis, and P x2supp(µi) wj(x) = bj for all i,
along with wj(x) 0. In particular, this implies that P
x2Sj wj(x) = kbj for all j = 1, . . . , n.
For arbitrary p, such a concise description of ⌫j is not possible. Therefore an alternate, but equivalent, way to characterize the distribution ⌫ is to just define the sets Sj and weight functions wj(·) for 1  j  n. This motivates the following definitions.
Definition 2.1. A solution (S,w) = (S1, . . . , Sn, w1, . . . , wj) is a valid partition as described previously (meaning that these partitions come from the optimal coupling between each µi to a fixed ⌫), along with the corresponding weight functions wj(·).
Definition 2.2. Let (S,w) be a solution. The cost of this solution, denoted costp(S), is the value of the objective (1) when we reconstruct ⌫ from S and w and evaluate (1):
costp(S) = min ⌫
1
k
nX
j=1
X
x2Sj
wj(x)kx ⌫ j k p.
Similarly for a projection ⇡, costp(⇡S) denotes the value of the objective (1) when we first project each of the distributions to Rm using ⇡, then compute e⌫ using the original weights wj:
costp(⇡S) = min e⌫
1
k
nX
j=1
X
x2Sj
wj(x)k⇡(x) e⌫jkp.
Note that each e⌫j 2 Rm. We suppress the dependence of the cost on w for notational convenience.
For the case of p = 2, we can further massage the value of ⌫j in (2). Let x̄ denote the weighted average of points in Sj (given by (3)). From our discussion above, we know that ⌫j = x̄. After some standard algebraic manipulation, we can show that
P x2Sj wj(x)kx ⌫ j k 2 = P x2Sj wj(x)kxk 2 kbjkx̄k2
and P
x,y2Sj wj(x)wj(y)kx yk 2 = 2kbj ⇣P x2Sj wj(x)kxk 2 kbjkx̄k2 ⌘ . Combining these
equations yields the following for the p = 2 objective.
1
2kbj
0 @ X
x,y2Sj
wj(x)wj(y)kx yk 2
1 A = X
x2Sj
wj(x)kx ⌫ j k 2. (4)
Dimension reduction. In this paper we are concerned with dimensionality reduction maps ⇡ : Rd ! Rm that are JL projections, i.e., any dimensionality reduction map that satisfies the condition of the JL lemma. This includes random Gaussian and sub-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, · · · , µk, solves for some approximate or exact ⌫ minimizing the objective (1). We can combine any such A with dimensionality reduction by first projecting the point masses of the µi down to Rm for some m < d and using A to compute some barycenter e⌫ in Rm. Then, we can consider the solution (S,w) induced by e⌫ (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 1 since we are given Sj and wj(·). For p = 2 (which is the most common case), ⌫j 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, if algorithm A 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 A is usually 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 m ⌧ d, we get significant savings since T (n, k,m) ⌧ T (n, k, d).
Algorithm 1 Using dimensionality reduction with any algorithm A for computing WB
Require: k discrete distributions µ1, · · · , µk with point masses in dimension Rd, projection dimension m, algorithm A
1: Project the point masses of each distribution µi to dimension Rm using a JL projection 2: Use algorithm A to solve (or approximately solve) the Wasserstein barycenter problem in Rm to
get a distribution e⌫ //e⌫ is a discrete distribution in Rm 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 2 S do 5: Solve for ⌫j minimizing
P x2Sj wj(x)kx ⌫ j k p //This is a convex program for
p 1. For p = 2, ⌫j is just the weighted average of points in Sj. 6: end for 7: Output the distribution ⌫ supported on ⌫j , and where ⌫j has the same weight as e⌫j
3 Reduction to O(log(nk)) Dimensions
We first show that it suffices to project the point masses of the input distribution into O(log(nk)) dimensions and guarantee that the cost of any solution is preserved. Note that our results hold simultaneously for all solutions. We first state the p = 2 case. Theorem 3.1. Consider a JL projection ⇡ from Rd to Rm for m = O(log(nk/ )/"2). Then
P cost2(⇡S) 2 [(1 ")2 · cost2(S) , (1 + ")2 · cost2(S)] for all solutions S
1 .
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 1 .
A decomposition similar to (4) does not exist for p 6= 2. To prove an analogous theorem for p 6= 2, we need the following Lipschitz extension theorem which roughly allows us to “invert” a dimensionality reduction map. Theorem 3.2 (Kirszbraun Theorem [Kir34]). For any D ⇢ Rm, let f : D ! Rd be an L-Lipschitz function. Then there exists some extension ef : Rm ! Rd of f to the entirety of Rm such that f(x) = ef(x) for all x 2 D and ef is also L-Lipschitz.
The Kirszbraun theorem allows us to prove Theorem 3.1 for general p with a dimension bound of m = O(log(nk/ )p2/"2) (see Theorem A.1 in Supplementary Section A).
The overview for the proof strategy for the general p 6= 2 case is as follows. First suppose that all the pairwise distances between the support points of all the distributions are preserved under the projection map up to multiplicative error 1± ". This event happens with probability at least 1 . We then consider the map f : Rm ! Rd that maps each of the projected points to its original counterpart in Rd. Note that the map is from the smaller dimension m to the larger dimension d. On the support points, we know that f is (1 + ")-Lipschitz by our assumption above.
Now if the projection caused the cost of ⇡S to decrease significantly, then using the Kriszbraun theorem, one could “lift” the corresponding barycenter e⌫ from the projected dimension to the original dimension using the extension map ef . Then since ef is Lipschitz, this lifted barycenter ef(e⌫) 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 ⌫ given in Section 2. Note that the exact description of ef does not matter for the analysis, just that such a map exists. A complete, rigorous proof can be found in the supplementary section.

heading: 4 Optimal Dimensionality Reduction
text: We now present our main theorem which improves the guarantees of Theorems 3.1 and A.1.
Theorem 4.1. Let µ1, . . . , µk be discrete probability distributions on Rd such that |supp(µi)|  poly(n) for all i. Let d 1, ", 2 (0, 1), and p 1. Let ⇡d,m : Rd ! Rm be a family of random JL maps with m = O ⇣ p4
"2 log n "
⌘ . Then we have,
P (costp(⇡S) 2 [(1 ") · costp(S) , (1 + ") · costp(S)] for all solutions S) 1 .
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.1. The key bottleneck is that when we project down to the m specified in Theorem 4.1, a large number of pairwise distances between the support points of the k distributions can be distorted (since we are projecting to a dimension smaller than O(log(nk))). Therefore, the Kirszbraun theorem cannot apply as the map f described in the proof strategy of Theorem A.1 is no longer Lipschitz on the support points.
To overcome this barrier, we generalize an approach of [MMR19], who achieved the optimal dimensionality bounds for k-means clustering beyond the naïve JL bound by defining a 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+")-factor under the random projection map ⇡. They show that the distortion graph is everywhere sparse, i.e., each vertex has small expected degree in the distortion graph, which implies a “robust” Kirszbraun theorem (for their particular problem of k-means clustering). Namely, there exists an extension map ef : Rd ! Rm and a specific point v 2 Rm in the projected space such that a large fraction of the distances from the pre-image ef 1(v) to the input points in Rd are preserved. Moreover, the input points whose distance to ef 1(v) is not preserved can be shown to contribute small error 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 we 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.

heading: 4.1 Dimensionality Reduction Lower Bounds
text: 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 + ")-approximation to the optimal Wasserstein barycenter requires dimension ⌦(log n/"2).
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 in the original dimension if the projection dimension is smaller than log n.
Theorem 4.3. There exists point sets A,B ⇢ Rd with |A| = |B| = n and matching cost M between them, such that if randomly projected down to m = o(log n) dimensions using an appropriately scaled Gaussian random matrix, the pull back cost of the optimal matching in Rm is at least !(M).
In addition, we prove a 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 ⇢ Rd with |A| = |B| = n and matching cost M between them, such that if randomly projected down to m = o(log n) dimensions using an appropriately scaled Gaussian random matrix, the cost of optimal matching in Rm is o(M) with probability at least 2/3.
See Supplementary Section D for full proofs.

heading: 5 Coresets
text: 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 on n. We first define our notion of coresets. Definition 5.1 (Coreset). Fix p 1. Let C and M be two sets of distributions in Rd where all distributions consist of poly(n) point masses. C is called an "-corset for the set of distributions M if there exist weights wc for c 2 C such that for all distributions ⌫ of support size at most n, it holds that
(1 ") X
c2C wc W (c, ⌫)
p 
1
|M |
X
µ2M W (µ, ⌫)p  (1 + ")
X c2C wc W (c, ⌫) p.
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 Rd, each supported on at most poly(n) point masses. There exists a weighted subset K ✓ M of size poly(n, d)/"2 that satisfies Definition 5.1 for p = O(1).
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 to 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 in our proofs and instead focus on showing that k can be reduced to 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 set N of all possible barycenter distributions ⌫ with support size at most n. The sensitivity of a distribution µ 2 M is defined as
(µ) = sup ⌫2N
W (µ, ⌫)p
1 |M | P µ02M W (µ 0, ⌫)p .
The total sensitivity is defined as S = 1|M | P µ2M (µ).
To see why such a 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 it is clear that we must sample µ with a higher probability if we wish to satisfy the definition of a 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, we utilize the notion of psuedo-dimension. Definition 5.4 (Pseudo-Dimension, Definition 9 [LFKF18]). Let X be a ground set and F be a set of functions from X to the interval [0, 1]. Fix a set S = {x1, · · · , xn} ⇢ X , a set of reals numbers R = {r1, · · · , rn} with ri 2 [0, 1] and a function f 2 F . The set Sf = {xi 2 S | f(xi) ri} is called the induced subset of S formed by f and R. The set S with associated values R is shattered by F if |{Sf | f 2 F}| = 2n. The pseudo-dimension of F is the cardinality of the largest shattered subset of X (or 1).
The following theorem provides a formal connection between the size of coresets and the notion of sensitivity and psuedo-dimension. Note that the statement of the theorem is more general and applies 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 " > 0 and 2 (0, 1). Let s : M ! R 0 denote any upper bound function on the sensitivity (·) defined in Definition 5.3 and let S = 1|M | P µ2M s(µ). Consider a set K of |K| samples of M with replacement where each distribution µ 2 M is sampled with probability q(µ) = s(µ)/(|M | · S) and each sampled point is assigned the weight 1/(|M | · |K| · q(µ)). Let F
denote the set of functions
F =
( W (·, ⌫)p
Sq(·) P µ2M W (µ, ⌫) p | ⌫ 2 N
)
where N is the set of all possible barycenter distributions with support size at most n. Let d0 denote the pseudo-dimension of F . Then the set K (along with the associated weights) satisfies Definition 5.1 with probability at least 1 if
|K| cS
"2
✓ d0 logS + log 1 ◆
where c > 0 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 is a metric. The latter requires us to use 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.

heading: 6 Other Theoretical Results
text: We now present some additional theoretical results pertaining to 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 µ1, . . . , µk 2 Rd with support size at most n, consider the problem of computing the Wasserstein barycenter with support size at most n for the p = 2 objective. There exists a matrix A 2 Rnk⇥d and a set S of rank n orthogonal projection matrices in Rnk⇥nk such that the first problem is equivalent to computing
P⇤ = argmin P2S kA PAk2F .
The proof of Theorem 6.1 is given in Section E.
We also prove the following NP hardness result in 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 a multiplicative factor 1.0013.

heading: 7 Experiments
text: 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 k = 2 distributions by splitting the images facing to the “left” versus the ones facing “right.” This results in ⇠ 350 uniform point masses per distribution.
MNIST dataset: We subsample 104 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 in R784. Experimental setup. We project our datasets in dimensions d ranging from d = 2 to d = 30 and compute the Wasserstein barycenter for p = 2. For FACES, we limit the support size of the barycenter to be at most 5 points in R4096 (since the barycenter should intuitively return an “interpolation” between the left and right facing faces, it should not be 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 in the higher dimension (see Algorithm 1) against the Wasserstein barycenter found in the higher dimension.
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 to d = 30 only increases the cost of the solution by 5%. This is 1-2 orders of magnitude smaller than from the original dimensions of 784 and 4096 for MNIST and FACES respectively. The average time taken to run the Wasserstein barycenter computation algorithm in d = 30 was 73% and 9% of the time taken to run in the full dimensions respectively.
Coreset experiments. Our coreset result reduces the number of distributions k through sensitivity (importance) sampling. We created a synthetic dataset with large k but small n and d to emphasize the advantage of sensitivity sampling over uniform sampling. We have k = 50, 000 distributions that each consists of a single point mass in R. The first k 1 distributions are all supported at the origin while one distribution is supported at x = k. We consider the p = 2 case and limit the support size of the barycenter to also be 1. Let costorig(⌫) denote the cost of ⌫ on the original objective (1) and let costcore(⌫) the cost of (1) when evaluated on a coreset. We record the relative error |costcore(⌫) costorig(⌫)|/|costorig(⌫)| evaluated at ⌫ = x, i.e. a single unit point mass at x, for x = 0, 1, 10. We then average the results across 10 trials each. As x (the point on which the query distribution is 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.

heading: Acknowledgments
text: Sandeep Silwal was supported in part by a NSF Graduate Research Fellowship Program. Samson Zhou was supported by a Simons Investigator Award of David P. Woodruff.

Review summaries:
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.

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. 

This paper present a method to reduce the dimension of k discrete measures defined on (at most) n support points in a real space of dimension d, by an O(log(n)) factor, while the optimum Wasserstein barycenter value maintains a bounded error from the optimal cost.  The main results of this work are given in Theorem 4.1 and 4.2. The paper presents the main plan of the technical proofs, but all the details of the proof are deferred to the online supplement of the paper. The two key ingredients of the proof are Johnson-Lindenstrauss (JL) projection results and the Kirszbraun theorem.  The paper contains also basic computational results on two standard datasets, the FACES and the MNIST datasets, which show the impact of the dimensionality reduction on the optimal value of the barycenter problem. 

This paper introduces dimensionality reduction techniques and coresets for Wasserstein barycenters with fixed support size n. More precisely, motivated by the fact that exact algorithms have a complexity that depend exponentially on the dimension d, the authors first show that for p-Wasserstein barycenters, a random projection to a space of dimension $O(p^4/\epsilon^2 \log n / \epsilon)$ that preserves the cost of the barycenter within a $(1\pm \epsilon)$ factor can be obtained with high probability, independently of the original dimension d, the size of the distributions, and their number. Second, the authors show that for a given set of mesures $\mu_1, …, \mu_M$, there exists a subset of $\mathrm{poly}(d, n) /\epsilon^2$ size with corresponding weights such that for any target measure $\nu$, the original cost lies within $(1 \pm \epsilon)$ factor w.r.t. the weighted cost on this coreset.