Dimensionality Reduction for Wasserstein Barycenter

1 Introduction . 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 . 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 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 . 1.2 Related Work . [ 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 .

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