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