Stochastic Approximation versus Sample Average Approximation for population Wasserstein barycenters
Weierstrass Institute for Applied Analysis and Stochastics（魏尔斯特拉斯应用分析与随机研究所）；Moscow Institute of Physics and Technology（莫斯科物理技术学院）；Institute for Information Transmission Problems（信息传输问题研究所）
In the machine learning and optimization community, there are two main approaches for the convex risk minimization problem, namely, the Stochastic Approximation (SA) and the Sample Average Approximation (SAA). In terms of oracle complexity (required number of stochastic gradient evaluations), both approaches are considered equivalent on average (up to a logarithmic factor). The total complexity depends on the specific problem, however, starting from work \cite{nemirovski2009robust} it was generally accepted that the SA is better than the SAA. % Nevertheless, in case of large-scale problems SA may run out of memory as storing all data on one machine and organizing online access to it can be impossible without communications with other machines. SAA in contradistinction to SA allows parallel/distributed calculations. We show that for the Wasserstein barycenter problem this superiority can be inverted. We provide a detailed comparison by stating the complexity bounds for the SA and the SAA implementations calculating barycenters defined with respect to optimal transport distances and entropy-regularized optimal transport distances. As a byproduct, we also construct confidence intervals for the barycenter defined with respect to entropy-regularized optimal transport distances in the $\ell_2$-norm. The preliminary results are derived for a general convex optimization problem given by the expectation in order to have other applications besides the Wasserstein barycenter problem.
We present a stochastic algorithm to compute the barycenter of a set of probability distributions under the Wasserstein metric from optimal transport. Unlike previous approaches, our method extends to continuous input distributions and allows the support of the barycenter to be adjusted in each iteration. We tackle the problem without regularization, allowing us to recover a sharp output whose support is contained within the support of the true barycenter. We give examples where our algorithm recovers a more meaningful barycenter than previous work. Our method is versatile and can be extended …
Georgia Institute of Technology（佐治亚理工学院）；University of California, Irvine（加利福尼亚大学尔湾分校）
Wasserstein Barycenter is a principled approach to represent the weighted mean of a given set of probability distributions, utilizing the geometry induced by optimal transport. In this work, we present a novel scalable algorithm to approximate the Wasserstein Barycenters aiming at high-dimensional applications in machine learning. Our proposed algorithm is based on the Kantorovich dual formulation of the Wasserstein-2 distance as well as a recent neural network architecture, input convex neural network, that is known to parametrize convex functions. The distinguishing features of our method ar…
In many applications in statistics and machine learning, the availability of data samples from multiple possibly heterogeneous sources has become increasingly prevalent. On the other hand, in distributionally robust optimization, we seek data-driven decisions which perform well under the most adverse distribution from a nominal distribution constructed from data samples within a certain discrepancy of probability distributions. However, it remains unclear how to achieve such distributional robustness in model learning and estimation when data samples from multiple sources are available. In thi…
Optimal transport is a notoriously difficult problem to solve numerically, with current approaches often remaining intractable for very large scale applications such as those encountered in machine learning. Wasserstein barycenters -- the problem of finding measures in-between given input measures in the optimal transport sense -- is even more computationally demanding as it requires to solve an optimization problem involving optimal transport distances. By training a deep convolutional neural network, we improve by a factor of 60 the computational speed of Wasserstein barycenters over the fas…
MIT CSAIL（麻省理工学院计算机科学与人工智能实验室）；IBM Research（IBM研究院）；MIT-IBM Watson AI Lab（麻省理工学院-IBM沃森人工智能实验室）
Wasserstein barycenters provide a geometrically meaningful way to aggregate probability distributions, built on the theory of optimal transport. They are difficult to compute in practice, however, leading previous work to restrict their supports to finite sets of points. Leveraging a new dual formulation for the regularized Wasserstein barycenter problem, we introduce a stochastic algorithm that constructs a continuous approximation of the barycenter. We establish strong duality and use the corresponding primal-dual relationship to parametrize the barycenter implicitly using the dual potential…
IBM Research（IBM研究院）；MIT-IBM Watson AI Lab（麻省理工学院-IBM沃森人工智能实验室）
In this paper we propose to perform model ensembling in a multiclass or a multilabel learning setting using Wasserstein (W.) barycenters. Optimal transport metrics, such as the Wasserstein distance, allow incorporating semantic side information such as word embeddings. Using W. barycenters to find the consensus between models allows us to balance confidence and semantics in finding the agreement between the models. We show applications of Wasserstein ensembling in attribute-based classification, multilabel learning and image captioning generation. These results show that the W. ensembling is a…
College of Computer Science and Technology, Jilin University（吉林大学计算机科学与技术学院）；Key Laboratory of Symbolic Computation and Knowledge Engineering of Ministry of Education（教育部符号计算与知识工程重点实验室）
Wasserstein barycenter, built on the theory of optimal transport, provides a powerful framework to aggregate probability distributions, and it has increasingly attracted great attention within the machine learning community. However, it suffers from severe computational burden, especially for high dimensional and continuous settings. To this end, we develop a novel continuous approximation method for the Wasserstein barycenters problem given sample access to the input distributions. The basic idea is to introduce a variational distribution as the approximation of the true continuous barycenter…
Skolkovo Institute of Science and Technology（斯科尔科沃科学技术研究院）；Massachusetts Institute of Technology（麻省理工学院）
Wasserstein barycenters provide a geometric notion of the weighted average of probability measures based on optimal transport. In this paper, we present a scalable algorithm to compute Wasserstein-2 barycenters given sample access to the input measures, which are not restricted to being discrete. While past approaches rely on entropic or quadratic regularization, we employ input convex neural networks and cycle-consistency regularization to avoid introducing bias. As a result, our approach does not resort to minimax optimization. We provide theoretical analysis on error bounds as well as empir…
Efficiently aggregating data from different sources is a challenging problem, particularly when samples from each source are distributed differently. These differences can be inherent to the inference task or present for other reasons: sensors in a sensor network may be placed far apart, affecting their individual measurements. Conversely, it is computationally advantageous to split Bayesian inference tasks across subsets of data, but data need not be identically distributed across subsets. One principled way to fuse probability distributions is via the lens of optimal transport: the Wasserste…
Kyoto University（京都大学）；RIKEN AIP（理化学研究所先进智能项目中心）；Facebook AI Research（脸书人工智能研究院）；SliceX AI（SliceX人工智能）
The Wasserstein barycenter has been widely studied in various fields, including natural language processing, and computer vision. However, it requires a high computational cost to solve the Wasserstein barycenter problem because the computation of the Wasserstein distance requires a quadratic time with respect to the number of supports. By contrast, the Wasserstein distance on a tree, called the tree-Wasserstein distance, can be computed in linear time and allows for the fast comparison of a large number of distributions. In this study, we propose a barycenter under the tree-Wasserstein distan…
Computing Wasserstein barycenters (a.k.a. Optimal Transport barycenters) is a fundamental problem in geometry which has recently attracted considerable attention due to many applications in data science. While there exist polynomial-time algorithms in any fixed dimension, all known running times suffer exponentially in the dimension. It is an open question whether this exponential dependence is improvable to a polynomial dependence. This paper proves that unless P=NP, the answer is no. This uncovers a "curse of dimensionality" for Wasserstein barycenter computation which does not occur for Opt…
Télécom Paris（巴黎电信学院）；Institut Polytechnique de Paris（巴黎理工学院）；Ecole Polytechnique（巴黎综合理工学院）；Université de Toulouse（图卢兹大学）；INRAE（法国国家农业食品与环境研究院）；Aalto University（阿尔托大学）
This paper introduces a novel and generic framework to solve the flagship task of supervised labeled graph prediction by leveraging Optimal Transport tools. We formulate the problem as regression with the Fused Gromov-Wasserstein (FGW) loss and propose a predictive model relying on a FGW barycenter whose weights depend on inputs. First we introduce a non-parametric estimator based on kernel ridge regression for which theoretical results such as consistency and excess risk bound are proved. Next we propose an interpretable parametric model where the barycenter weights are modeled with a neural …
University of Vienna（维也纳大学）；Universidad de Chile（智利大学）；NoiseGrasp SpA（NoiseGrasp公司）
We introduce and study a novel model-selection strategy for Bayesian learning, based on optimal transport, along with its associated predictive posterior law: the Wasserstein population barycenter of the posterior law over models. We first show how this estimator, termed Bayesian Wasserstein barycenter (BWB), arises naturally in a general, parameter-free Bayesian model-selection framework, when the considered Bayesian risk is the Wasserstein distance. Examples are given, illustrating how the BWB extends some classic parametric and non-parametric selection strategies. Furthermore, we also provi…
We provide theoretical complexity analysis for new algorithms to compute the optimal transport (OT) distance between two discrete probability distributions, and demonstrate their favorable practical performance over state-of-art primal-dual algorithms and their capability in solving other problems in large-scale, such as the Wasserstein barycenter problem for multiple probability distributions. First, we introduce the \emph{accelerated primal-dual randomized coordinate descent} (APDRCD) algorithm for computing the OT distance. We provide its complexity upper bound $\bigOtil(\frac{n^{5/2}}{\var…
Entropy regularization in optimal transport (OT) has been the driver of many recent interests for Wasserstein metrics and barycenters in machine learning. It allows to keep the appealing geometrical properties of the unregularized Wasserstein distance while having a significantly lower complexity thanks to Sinkhorn's algorithm. However, entropy brings some inherent smoothing bias, resulting for example in blurred barycenters. This side effect has prompted an increasing temptation in the community to settle for a slower algorithm such as log-domain stabilized Sinkhorn which breaks the parallel …
In this work we introduce the concept of Bures-Wasserstein barycenter $Q_*$, that is essentially a Fréchet mean of some distribution $\mathbb{P}$ supported on a subspace of positive semi-definite Hermitian operators $\mathbb{H}_{+}(d)$. We allow a barycenter to be restricted to some affine subspace of $\mathbb{H}_{+}(d)$ and provide conditions ensuring its existence and uniqueness. We also investigate convergence and concentration properties of an empirical counterpart of $Q_*$ in both Frobenius norm and Bures-Wasserstein distance, and explain, how obtained results are connected to optimal tra…
Computationally solving multi-marginal optimal transport (MOT) with squared Euclidean costs for $N$ discrete probability measures has recently attracted considerable attention, in part because of the correspondence of its solutions with Wasserstein-$2$ barycenters, which have many applications in data science. In general, this problem is NP-hard, calling for practical approximative algorithms. While entropic regularization has been successfully applied to approximate Wasserstein barycenters, this loses the sparsity of the optimal solution, making it difficult to solve the MOT problem directly …
University of Wisconsin-Madison（威斯康星大学麦迪逊分校）；Google（谷歌）
Based on the concepts of Wasserstein barycenter (WB) and Gromov-Wasserstein barycenter (GWB), we propose a unified mathematical framework for neural network (NN) model fusion and utilize it to reveal new insights about the linear mode connectivity of SGD solutions. In our framework, the fusion occurs in a layer-wise manner and builds on an interpretation of a node in a network as a function of the layer preceding it. The versatility of our mathematical framework allows us to talk about model fusion and linear mode connectivity for a broad class of NNs, including fully connected NN, CNN, ResNet…
As interest in graph data has grown in recent years, the computation of various geometric tools has become essential. In some area such as mesh processing, they often rely on the computation of geodesics and shortest paths in discretized manifolds. A recent example of such a tool is the computation of Wasserstein barycenters (WB), a very general notion of barycenters derived from the theory of Optimal Transport, and their entropic-regularized variant. In this paper, we examine how WBs on discretized meshes relate to the geometry of the underlying manifold. We first provide a generic stability …
We propose a study of a distribution registration model for general deformation functions. In this framework, we provide estimators of the deformations as well as a goodness of fit test of the model. For this, we consider a criterion which studies the Fr{é}chet mean (or barycenter) of the warped distributions whose study enables to make inference on the model. In particular we obtain the asymptotic distribution and a bootstrap procedure for the Wasserstein variation.
Université Paris-Saclay（巴黎萨克雷大学）；CNRS（法国国家科学研究中心）；Inria（法国国家信息与自动化研究所）；Université Gustave Eiffel（古斯塔夫·埃菲尔大学）；University of Genoa（热那亚大学）；Istituto Italiano di Tecnologia（意大利技术研究院）；University College London（伦敦大学学院）
We study the problem of learning a real-valued function that satisfies the Demographic Parity constraint. It demands the distribution of the predicted output to be independent of the sensitive attribute. We consider the case that the sensitive attribute is available for prediction. We establish a connection between fair regression and optimal transport theory, based on which we derive a close form expression for the optimal fair predictor. Specifically, we show that the distribution of this optimum is the Wasserstein barycenter of the distributions induced by the standard regression function o…
IRIT, Université de Toulouse, CNRS（法国国家科学研究中心、图卢兹大学信息与电信研究所）；IDIA & CMM, Universidad de Chile（智利大学数据与人工智能研究所、智利大学数学建模中心）；CMM, Universidad de Chile（智利大学数学建模中心）
We introduce weak barycenters of a family of probability distributions, based on the recently developed notion of optimal weak transport of mass by Gozlanet al. (2017) and Backhoff-Veraguas et al. (2020). We provide a theoretical analysis of this object and discuss its interpretation in the light of convex ordering between probability measures. In particular, we show that, rather than averaging the input distributions in a geometric way (as the Wasserstein barycenter based on classic optimal transport does) weak barycenters extract common geometric information shared by all the input distribut…
This paper is focused on the statistical analysis of probability measures $ν_{1},\ldots,ν_{n}$ on $\mathbb{R}$ that can be viewed as independent realizations of an underlying stochastic process. We consider the situation of practical importance where the random measures $ν_{i}$ are absolutely continuous with densities $f_{i}$ that are not directly observable. In this case, instead of the densities, we have access to datasets of real random variables $(X_{i,j})_{1 \leq i \leq n; \; 1 \leq j \leq p_{i} }$ organized in the form of $n$ experimental units, such that $X_{i,1},\ldots,X_{i,p_{i}}$ are…
Université Paris-Saclay（巴黎萨克雷大学）；CNRS（法国国家科学研究中心）；Laboratoire de mathématiques d’Orsay（奥赛数学实验室）；Università di Genova（热那亚大学）
This work provides several fundamental characterizations of the optimal classification function under the demographic parity constraint. In the awareness framework, akin to the classical unconstrained classification case, we show that maximizing accuracy under this fairness constraint is equivalent to solving a corresponding regression problem followed by thresholding at level $1/2$. We extend this result to linear-fractional classification measures (e.g., ${\rm F}$-score, AM measure, balanced accuracy, etc.), highlighting the fundamental role played by the regression problem in this framework…
IBM Research（IBM研究院）；IBM T.J Watson Research Center（IBM托马斯·J·沃森研究中心）
We propose Gaussian optimal transport for Image style transfer in an Encoder/Decoder framework. Optimal transport for Gaussian measures has closed forms Monge mappings from source to target distributions. Moreover interpolates between a content and a style image can be seen as geodesics in the Wasserstein Geometry. Using this insight, we show how to mix different target styles , using Wasserstein barycenter of Gaussian measures. Since Gaussians are closed under Wasserstein barycenter, this allows us a simple style transfer and style mixing and interpolation. Moreover we show how mixing differe…
Wasserstein barycenters and variance-like criteria based on the Wasserstein distance are used in many problems to analyze the homogeneity of collections of distributions and structural relationships between the observations. We propose the estimation of the quantiles of the empirical process of Wasserstein's variation using a bootstrap procedure. We then use these results for statistical inference on a distribution registration model for general deformation functions. The tests are based on the variance of the distributions with respect to their Wasserstein's barycenters for which we prove cen…
INRIA（法国国家信息与自动化研究所）；Ecole Normale Supérieure（高等师范学校）；PSL Research University（巴黎文理研究大学）；Kayrros SAS（凯罗斯公司）；CNRS（法国国家科学研究中心）
Hyperspectral satellite images report greenhouse gas concentrations worldwide on a daily basis. While taking simple averages of these images over time produces a rough estimate of relative emission rates, atmospheric transport means that simple averages fail to pinpoint the source of these emissions. We propose using Wasserstein barycenters coupled with weather data to average gas concentration data sets and better concentrate the mass around significant sources.
HSE University（俄罗斯高等经济大学）；Moscow Institute of Physics and Technology（莫斯科物理技术学院）；Institute for Information Transmission Problems（信息传输问题研究所）；Weierstrass Institute for Applied Analysis and Stochastics（魏尔斯特拉斯应用分析与随机研究所）
We consider the population Wasserstein barycenter problem for random probability measures supported on a finite set of points and generated by an online stream of data. This leads to a complicated stochastic optimization problem where the objective is given as an expectation of a function given as a solution to a random optimization problem. We employ the structure of the problem and obtain a convex-concave stochastic saddle-point reformulation of this problem. In the setting when the distribution of random probability measures is discrete, we propose a stochastic optimization algorithm and es…
In this paper, we focus on computational aspects of the Wasserstein barycenter problem. We propose two algorithms to compute Wasserstein barycenters of $m$ discrete measures of size $n$ with accuracy $\e$. The first algorithm, based on mirror prox with a specific norm, meets the complexity of celebrated accelerated iterative Bregman projections (IBP), namely $\widetilde O(mn^2\sqrt n/\e)$, however, with no limitations in contrast to the (accelerated) IBP, which is numerically unstable under small regularization parameter. The second algorithm, based on area-convexity and dual extrapolation, im…
Wasserstein Barycenter Gaussian Process based Bayesian Optimization
Department of Economics Management and Statistics, University of Milano-Bicocca（经济管理统计系，米兰-比科卡大学）；Department of Computer Science Systems and Communication, University of Milano-Bicocca（计算机科学系统与通信系，米兰-比科卡大学）
Wasserstein barycenters correspond to optimal solutions of transportation problems for several marginals, and as such have a wide range of applications ranging from economics to statistics and computer science. When the marginal probability measures are absolutely continuous (or vanish on small sets) the theory of Wasserstein barycenters is well-developed (see the seminal paper [1]). However, exact continuous computation of Wasserstein barycenters in this setting is tractable in only a small number of specialized cases. Moreover, in many applications data is given as a set of probability measu…
We consider the problem of computing a Wasserstein barycenter for a set of discrete probability distributions with finite supports, which finds many applications in areas such as statistics, machine learning and image processing. When the support points of the barycenter are pre-specified, this problem can be modeled as a linear programming (LP) problem whose size can be extremely large. To handle this large-scale LP, we analyse the structure of its dual problem, which is conceivably more tractable and can be reformulated as a well-structured convex problem with 3 kinds of block variables and …
Wasserstein barycenter is the centroid of a collection of discrete probability distributions which minimizes the average of the $\ell_2$-Wasserstein distance. This paper focuses on the computation of Wasserstein barycenters under the case where the support points are free, which is known to be a severe bottleneck in the D2-clustering due to the large-scale and nonconvexity. We develop an inexact proximal alternating minimization (iPAM) method for computing an approximate Wasserstein barycenter, and provide its global convergence analysis. This method can achieve a good accuracy with a reduced …
We propose a new \cu{class-optimal} algorithm for the distributed computation of Wasserstein Barycenters over networks. Assuming that each node in a graph has a probability distribution, we prove that every node can reach the barycenter of all distributions held in the network by using local interactions compliant with the topology of the graph. We provide an estimate for the minimum number of communication rounds required for the proposed method to achieve arbitrary relative precision both in the optimality of the solution and the consensus among all agents for undirected fixed networks.
We study the problem of the decentralized computation of entropy-regularized semi-discrete Wasserstein barycenters over a network. Building upon recent primal-dual approaches, we propose a sampling gradient quantization scheme that allows efficient communication and computation of approximate barycenters where the factor distributions are stored distributedly on arbitrary networks. The communication and algorithmic complexity of the proposed algorithm are shown, with explicit dependency on the size of the support, the number of distributions, and the desired accuracy. Numerical results validat…
Computing the Wasserstein barycenter of a set of probability measures under the optimal transport metric can quickly become prohibitive for traditional second-order algorithms, such as interior-point methods, as the support size of the measures increases. In this paper, we overcome the difficulty by developing a new adapted interior-point method that fully exploits the problem's special matrix structure to reduce the iteration complexity and speed up the Newton procedure. Different from regularization approaches, our method achieves a well-balanced tradeoff between accuracy and speed. A numeri…
