A novel notion of barycenter for probability distributions based on optimal weak mass transport

1 Introduction . Optimal transport ( OT ) [ 41 ] has had a tremendous impact in the machine learning ( ML ) community recently , as it provides meaningful and implementable distances between probability distributions [ 34 ] , thus advancing many aspects in the field , see e.g . [ 7 , 43 , 33 ] . The space of probability measures on Rd with finite second moment can be metrised with the Wasserstein-2 distance , the computation of which amounts to finding a transport plan that minimises the quadratic average cost of transporting mass from a source probability measure onto a target one . In this context , a natural method for averaging a finite family of probability measures is to compute their Fréchet mean , with respect to the Wasserstein-2 distance , which corresponds to the Wasserstein barycenter introduced in [ 1 ] . The goal of the present work is to explore theoretical features and potential applications to ML of barycenters of probability measures analogously defined in terms of optimal weak transport ( OWT , see [ 26 ] ) or more precisely quadratic barycentric transport costs . In a nutshell , for a source measure µ and a target measure ν , the OWT problem aims to transport mass so that the conditional spatial mean of target support points y , given their source support points x , is close to x in average . This amounts to finding an intermediate measure η , possibly more concentrated than ν in the sense of convex ordering of probability measures , which is close to µ with respect to the Wasserstein-2 distance . The main motivation of our work is to investigate the effect and meaning of combining a family of probability measures using OWT instead of OT . To that end , we will define the weak barycenter of 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . this family through an optimisation problem , and discuss some of its properties . Importantly , we will see that , rather than averaging the input distributions in a metric sense , solving a weak barycenter problem corresponds to finding probability measures that encode geometric or shape information shared across all of them . In fact , the weak barycenter problem will be interpreted as finding a latent random variable common to all the input distributions . Implications of this latent variable interpretation , in terms of robustness to outliers , will also be drawn in our work . A second motivation for our work is to develop and implement computational methods for weak barycenters , capitalising on the fact that the optimal weak coupling between any pair of distributions , with finite second moments , is always realised by a unique optimal map . This property is in sharp contrast to standard OT , where the absolute continuity with respect to the Lebesgue measure of the source or target measure is typically needed to grant the existence and uniqueness of a map—the so-called Monge map— realising the optimal coupling between them . This map is often required in different ways to compute Wasserstein barycenters ( see [ 5 ] , [ 44 ] or [ 35 ] ) . Similarly to the Wasserstein barycenter problem , we will develop a fixed-point formulation of the weak barycenter problem , based on OWT plans . This allows us to , following [ 5 , 44 ] , construct an iterative procedure to compute a weak barycenter for a finite family of distributions and analyse its convergence properties . We will also define and study the so-called weak population barycenters , that deal with a population of probability measures distributed according to a given law Q supported on the Wasserstein-2 space , as in [ 30 ] for the OT case . Extending ideas from [ 10 ] , we will then propose an iterative stochastic algorithm for online computation of the weak population barycenter , from a stream of probability measures sampled from Q . We will then provide numerical simulations using this proposed method , in order to illustrate the geometric meaning of the weak barycenter , and we will compare it with related objects obtained with standard OT or its entropy-regularised counterpart . Organisation of the paper . Sec . 2 documents the background on OWT and the assumptions underlying our work . Sec . 3 analyses the weak barycenter problem , interprets it in the light of convex ordering and a latent variable model and addresses the case of an infinite population of distributions . Sec . 4 introduces two algorithms for computing the weak barycenter in the finite or population settings . Sec . 5 and 6 present the experimental setting and validation of our proposal respectively . Lastly , Sec . 7 discusses our findings and future research questions . The Appendix contains all the proofs , additional details or our simulations and the code of our experiments . 2 Background : optimal weak transport and Wasserstein barycenter . The optimal transport ( OT ) problem [ 42 ] aims to find the lowest cost to transfer the mass from one probability measure onto another . Therefore , OT is a natural way to compare two probability distributions in terms of their geometric information . In particular , the Wasserstein-p distance Wp , associated with the Euclidean cost in Rd , metrises the space Pp ( Rd ) of probability measures on Rd with finite p-moment . Precisely , for µ , ν ∈ Pp ( Rd ) , Wp ( µ , ν ) = ( min π∈Π ( µ , ν ) ∫ Rd×Rd ‖x− y‖pdπ ( x , y ) ) 1/p , ( 1 ) where π is a transport plan between µ and ν , that is , an element of the set Π ( µ , ν ) of probability measures on the product space Rd × Rd with marginals µ and ν . For p = 2 and µ absolutely continuous ( a.c. ) , the unique optimal plan is concentrated on the graph of a measurable map called Monge map such that ν = T # µ , see eq . ( 14 ) in Appendix A.1 . Optimal weak transport . We consider here the optimal weak transport ( OWT ) problem introduced in [ 26 ] and in particular the special case of barycentric transport costs . The OWT problem is then defined for µ , ν ∈ P2 ( Rd ) by V ( µ|ν ) = inf π∈Π ( µ , ν ) ∫ Rd ‖x− ∫ Rd ydπx ( y ) ‖2dµ ( x ) , ( 2 ) where πx is the disintegration of the transport plan π with respect to the first marginal µ , i.e . π ( dxdy ) = πx ( dy ) µ ( dx ) . As our work strongly leans on OWT theory , we recall in Appendix A.2 , Th . 6 , that V is continuous with respect to the Wasserstein metric [ 9 ] . Additionally , the two following results from [ 8 ] ( stated for our specific setting ) lay the ground for our proposed weak barycenters . Theorem 1 ( [ 8 ] , Theorem 1.2 ) . The problem ( 2 ) admits a unique minimiser . This first result strongly differs from the classical OT setting , for which the uniqueness of an optimal transport plan is not guaranteed for arbitrary measures . The optimisation problem in Eq . ( 2 ) can also be reformulated thanks to the Brenier-Strassen theorem [ 25 ] , [ 8 ] , through the notion of convex ordering . We denote by η ≤c ν the convex order of measures , meaning that ∫ φdη ≤ ∫ φdν for any convex function φ that is nonnegative or integrable with respect to η + ν . By Strassen ’ s theorem [ 40 ] , two distributions are in convex order if and only if there exists a martingale coupling between them . The following theorem is a generalisation of the result originally proved in [ 25 ] , Th . 1.2 . Theorem 2 ( [ 8 ] , Theorem 1.4 ) . Let µ ∈ P2 ( Rd ) and ν ∈ P1 ( Rd ) . There exists a unique η∗ ≤c ν such that W 22 ( µ , η ∗ ) = inf η≤cν W 22 ( µ , η ) = V ( µ|ν ) . ( 3 ) Moreover , there exists a convex function ψ : Rd → R of classC1 with∇ψ being 1-Lipschitz , such that ∇ψ # µ = η∗ . Finally , the optimal coupling πµ , ν ∈ Π ( µ , ν ) verifies ∫ ydπµ , νx ( y ) = ∇ψ ( x ) µ-a.s . The measurable map , or barycentric projection , Sνµ ( x ) : = ∫ Rd ydπ µ , ν x ( y ) associated to the plan πµ , ν achieving the minimum in Eq . ( 2 ) is consequently uniquely defined and will be called optimal barycentric projection . From this notation , we can write the OWT cost in terms of an OT cost according to V ( µ|ν ) = W 22 ( µ , Sνµ # µ ) . We emphasise that Sνµ is directly related to the optimisation problem ( 2 ) , whereas applied works such as [ 38 , 36 ] make use of a barycentric projection constructed from a transport plan solving an OT problem between µ and ν ( often regularised ) as a substitute for the Monge map , which may not exist ( more details on Sνµ are displayed in Appendix A.3 ) . Last , let us note that OWT is somehow also related to the martingale OT problem developed in the stochastic finance community [ 13 , 2 , 27 ] , which puts the focus on the optimal transfer of mass between distributions assumed to be in convex order themselves . Wasserstein barycenter . The classical Wasserstein barycenter problem for a set of probability measures ν1 , . . . , νn ∈ P2 ( Rd ) with weights λ1 , . . . , λn in the simplex ( i.e . λi ≥ 0 and ∑n i=1 λi = 1 ) is defined [ 1 ] by arg min µ∈P2 ( Rd ) n∑ i=1 λiW 2 2 ( µ , νi ) . ( 4 ) The Wasserstein barycenter has been extensively studied both theoretically and numerically [ 30 , 44 , 5 , 15 ] . Regarding the numerical part , [ 39 ] focuses on the computation of Wasserstein barycenters for a fixed number of measures and a stream of observations per measure ; additionally , [ 32 ] proposed an entropy-regularised alternative via stochastic optimisation for computing the Wasserstein barycenter of a.c. distributions only from observations . Constrained by their assumption of a.c. , [ 44 ] computes the Wasserstein barycenter by smoothing the observed empirical distributions . Furthermore , [ 21 ] compares the complexity of both the sample Wasserstein barycenter and a stochastic approximation to estimate a population barycenter ( discrete measures and entropic regularisation ) . Finally , the authors of [ 3 ] recently proposed an algorithm to compute the barycenters in polynomial time . 3 Optimal weak transport barycenters and latent variable interpretation . 3.1 Definition and basic properties . In a similar fashion , based on the weak transport cost in Eq . ( 2 ) , we propose the following variant : Definition 1 . The set of weak barycenters of a finite family of measures { νi } i=1 , ... , n ∈ P2 ( Rd ) with weights { λi } i=1 , ... , n in the simplex is defined as arg min µ∈P2 ( Rd ) n∑ i=1 λiV ( µ|νi ) . ( 5 ) Thus , a weak barycenter averages , with respect to the Wasserstein metric , an optimally chosen set of probability measures { η1 , . . . , ηn } which are more concentrated than the corresponding νi , in the sense that ηi ≤c νi for each 1 ≤ i ≤ n. The existence of a solution is established as follows : Proposition 1 . The weak barycenter problem in Eq . ( 5 ) admits a minimiser µ ∈ P2 ( Rd ) . See Sec . B of the Appendix for the proof of the above Proposition ( which relies on Prokhorov ’ s theorem ) and all the proofs for this Section . Uniqueness is in general not granted : we next show that the set of solutions is indeed an interval , with respect to the partial order of convex ordering of probability measures . In the following , we denote by X and Yi random variables with respective laws µ and νi , for 1 ≤ i ≤ n , and δa the Dirac measure supported on a ∈ Rd . Lemma 1 . If µ is a weak barycenter of { νi } i=1 ... , n and µ′ ≤c µ , then µ′ also is a weak barycenter . In particular , the Dirac measure supported on Eµ ( X ) is always a weak barycenter . Moreover , a Dirac distribution δω̄ is a weak barycenter if and only if ω̄ = ∑n i=1 λiEνi ( Yi ) . A consequence of the above lemma is that for any weak barycenter µ , Eµ ( X ) = n∑ i=1 λiEνi ( Yi ) , ( 6 ) and the value of the weak barycenter problem is given by inf µ∈P2 ( Rd ) n∑ i=1 λiV ( µ|νi ) = n∑ i=1 λi‖E ( Yi ) ‖2 − ‖ n∑ i=1 λiE ( Yi ) ‖2 . ( 7 ) We can also derive the following characterisation on the set of weak barycenters : Proposition 2 . A measure µ ∈ P ( Rd ) is a weak barycenter of { νi } i=1 ... , n if and only if its mean satisfies ( 6 ) and µ̂ ≤c ν̂i holds for all 1 ≤ i ≤ n , where ν̂ denotes the centered version of a law ν . For instance , in the case of one dimensional Gaussian distributions νi = N ( m , σ2i ) , the set of weak barycenters includes { µ = N ( m , σ2 ) | 0 ≤ σ2 ≤ min1≤i≤n σ2i } . A natural question is whether a `` maximal ” weak barycenter exists , in the sense of convex ordering ( up to translation by the mean ) . For d = 1 , the answer is affirmative . When the means E ( Yi ) are equal , this follows from the complete lattice property of the set of probability measures with respect to the convex ordering ( see [ 29 ] ) ; the general case can then be reduced to the latter using Proposition 2 . For d ≥ 2 , this property is in general not true and the answer depends on the family { νi } i=1 ... , n. In the particular case of a.c. input measures , we can bound the distance between the Wasserstein and weak barycenters by the variances of the distributions ( νi ) 1≤i≤n . The barycenters are then closer the more concentrated each νi is . Lemma 2 . Let ν1 , . . . , νn ∈ P2 ( Rd ) be a.c. , at least one of them with bounded density . Let µ̄ and µ̃ respectively denote the weak and the Wasserstein barycenters . Then W 22 ( µ̄ , µ̃ ) ≤ 2 n∑ i=1 λi ( E‖Yi‖2 − ‖EYi‖2 ) .

Summary: This paper studies barycenters of probability distributions using the weak transport distance, and establishes a number of results which mirror the theory of standard optimal transport barycenters that has been established over the past 10 years. In particular, they investigate existence and uniqueness, the first order optimality criterion, as well as iterative methods to compute weak barycenters. They propose a proximal sub-routine for this iterative method, and evaluate the performance of weak barycenters on a number of examples, comparing with standard barycenters and Sinkhorn barycenters.