Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization

1 INTRODUCTION . Wasserstein barycenters have become popular due to their ability to represent the average of probability measures in a geometrically meaningful way . Techniques for computing Wasserstein barycenters have been successfully applied to many computational problems . In image processing , Wasserstein barycenters are used for color and style transfer ( Rabin et al. , 2014 ; Mroueh , 2019 ) , and texture synthesis ( Rabin et al. , 2011 ) . In geometry processing , shape interpolation can be done by computing barycenters ( Solomon et al. , 2015 ) . In online machine learning , barycenters are used for aggregating probabilistic predictions of experts ( Korotin et al. , 2019b ) . Within the context of Bayesian inference , the barycenter of subset posteriors converges to the full data posterior , thus enabling efficient computational methods based on finding the barycenters ( Srivastava et al. , 2015 ; 2018 ) . Fast and accurate barycenter algorithms exist for discrete distributions ( see Peyré et al . ( 2019 ) for a survey ) , while for continuous distributions the situation is more difficult and remains unexplored until recently ( Li et al. , 2020 ; Fan et al. , 2020 ; Cohen et al. , 2020 ) . The discrete methods scale poorly with the number of support points of the barycenter and thus can not approximate continuous barycenters well , especially in high dimensions . In this paper , we present a method to compute Wasserstein-2 barycenters of continuous distributions based on a novel regularized dual formulation where the convex potentials are parameterized by input convex neural networks ( Amos et al. , 2017 ) . Our algorithm is straightforward without introducing bias ( e.g . Li et al . ( 2020 ) ) or requiring minimax optimization ( e.g . Fan et al . ( 2020 ) ) . This is made possible by combining a new congruence regularizing term combined with cycle-consistency regularization ( Korotin et al. , 2019a ) . As we will show in the analysis , thanks to the properties of Wasserstein-2 distances , the gradients of the resulting convex potentials “ push ” the input distributions close to the true barycenter , allowing good approximation of the barycenter . 2 PRELIMINARIES . We denote the set of all Borel probability measures on RD with finite second moment by P2 ( RD ) . We use P2 , ac ( RD ) ⊂ P2 ( RD ) to denote the subset of all absolutely continuous measures ( w.r.t . the Lebesgue measure ) . Wasserstein-2 distance . For P , Q ∈ P2 ( RD ) , the Wasserstein-2 distance is defined by W22 ( P , Q ) def = min π∈Π ( P , Q ) ∫ RD×RD ‖x− y‖2 2 dπ ( x , y ) , ( 1 ) where Π ( P , Q ) is the set of probability measures on RD×RD whose marginals are P , Q , respectively . This definition is known as Kantorovich ’ s primal form of transport distance ( Kantorovitch , 1958 ) . The Wasserstein-2 distance W2 is well-studied in the theory of optimal transport ( Brenier , 1991 ; McCann et al. , 1995 ) . In particular , it has a dual formulation ( Villani , 2003 ) : W22 ( P , Q ) = ∫ RD ‖x‖2 2 dP ( x ) + ∫ RD ‖y‖2 2 dQ ( y ) − min ψ∈Conv [ ∫ RD ψ ( x ) dP ( x ) + ∫ RD ψ ( y ) dQ ( y ) ] , ( 2 ) where the minimum is taken over all the convex functions ( potentials ) ψ : RD → R ∪ { ∞ } , and ψ ( y ) = maxx∈RD ( 〈x , y〉−ψ ( x ) ) : RD → R∪ { ∞ } is the convex conjugate of ψ ( Fenchel , 1949 ) , which is also a convex function . The optimal potential ψ∗ is defined up to an additive constant . Brenier ( 1991 ) shows that if P does not give mass to sets of dimensions at most D − 1 , then the optimal plan π is uniquely determined by π∗ = [ idRD , T ∗ ] ] P , where T ∗ : RD → RD is the unique solution to the Monge ’ s problem T ∗ = arg min T ] P=Q ∫ RD ‖x− T ( x ) ‖2 2 dP ( x ) . ( 3 ) The connection between T ∗ and the dual formulation ( 2 ) is that T ∗ = ∇ψ∗ , where ψ∗ is the optimal solution of ( 2 ) . Additionally , if Q does not give mass to sets of dimensions at most D − 1 , then T ∗ is invertible and T ∗ ( x ) = ∇ψ∗ ( x ) = ( ∇ψ∗ ) −1 ( x ) , ( T ∗ ) −1 ( y ) = ∇ψ∗ ( y ) = ( ∇ψ∗ ) −1 ( y ) . In particular , the above discussion applies to the case where P , Q ∈ P2 , ac ( RD ) . Wasserstein-2 barycenter . Let P1 , . . . , PN ∈ P2 , ac ( RD ) . Then , their barycenter w.r.t . weights α1 , . . . , αN ( αn > 0 and ∑N n=1 αn = 1 ) is P def= arg min P∈P2 ( RD ) N∑ n=1 αnW22 ( Pn , P ) . ( 4 ) Throughout this paper , we assume that at least one of P1 , . . . , PN ∈ P2 , ac ( RD ) has bounded density . Under this assumption , P is unique and absolutely continuous , i.e. , P ∈ P2 , ac ( RD ) , and it has bounded density ( Agueh & Carlier , 2011 , Definition 3.6 & Theorem 5.1 ) . For n ∈ { 1 , 2 , . . . , N } , let ( ψ∗n , ψ∗n ) be the optimal pair of ( mutually ) conjugate potentials that transport Pn to P , i.e. , ∇ψ∗n ] Pn = P and∇ψ∗n ] P = Pn . Then { ψ∗n } satisfy N∑ n=1 αn∇ψ∗n ( x ) = x and N∑ n=1 αnψ∗n ( x ) = ‖x‖2 2 + c. ( 5 ) for all x ∈ RD ( Agueh & Carlier , 2011 ; Álvarez-Esteban et al. , 2016 ) . Since optimal potentials are defined up to a constant , for convenience , we set c = 0 . The condition ( 5 ) serves as the basis for our algorithm for computing Wasserstein-2 barycenters . We say that potentialsψ1 , . . . , ψN are congruent w.r.t . weights α1 , . . . , αn if their conjugate potentials satisfy ( 5 ) , i.e. , ∑D n=1 αnψn ( x ) = ‖x‖2 2 for all x ∈ RD . 3 RELATED WORK . Most algorithms in the field of computational optimal transport are designed for the discrete setting where the input distributions have finite support ; see the recent survey by Peyré et al . ( 2019 ) for discussion . A particular popular line of algorithms are based on entropic regularization that gives rise to the famous Sinkhorn iteration ( Cuturi , 2013 ; Cuturi & Doucet , 2014 ) . These methods are typically limited to a support of 105−106 points before the problem becomes computationally infeasible . Similarly , discrete barycenter methods ( Cuturi & Doucet , 2014 ) , particularly the ones that rely on a fixed support for the barycenter ( Dvurechenskii et al. , 2018 ; Staib et al. , 2017 ) , can not provide precise approximation of continuous barycenters in high dimensions , since a large number of samples is needed ; see experiments in Fan et al . ( 2020 , §4.3 ) for an example . Thus we focus on the existing literature in the continuous setting . Computation of Wasserstein-2 distances and maps . Genevay et al . ( 2016 ) demonstrate the possibility of computing Wasserstein distances given only sample access to the distributions by parameterizing the dual potentials as functions in the reproducing kernel Hilbert spaces . Based on this realization , Seguy et al . ( 2017 ) propose a similar method but use neural networks to parameterize the potentials , using entropic or L2 regularization w.r.t . P×Q to keep the potentials approximately conjugate . The transport map is recovered from optimized potentials via barycentric projection . As we note in §2 , W2 enjoys many useful theoretical properties . For example , the optimal potential ψ∗ is convex , and the corresponding optimal transport map is given by ∇ψ∗ . By exploiting these properties , Makkuva et al . ( 2019 ) propose a minimax optimization algorithm for recovering transport maps , using input convex neural networks ( ICNNs ) ( Amos et al. , 2017 ) to approximate the potentials . An alternative to entropic regularization is the cycle-consistency regularization proposed by Korotin et al . ( 2019a ) . It uses the property that the gradients of optimal dual potentials are inverses of each other . The imposed regularizer requires integration only over the marginal measures P and Q , instead of over P × Q as required by entropy-based alternatives . Their method converges faster than the minimax method since it does not have an inner optimization cycle . Xie et al . ( 2019 ) propose using two generative models with a shared latent space to implicitly compute the optimal transport correspondence between P and Q . Based on the obtained correspondence , the authors are able to compute the optimal transport distance between the distributions . Computation of Wasserstein-2 barycenters . A few recent techniques tackle the barycenter problem ( 4 ) using continuous rather than discrete approximations of the barycenter : • MEASURE-BASED ( GENERATIVE ) OPTIMIZATION : Problem ( 4 ) optimizes over probability measures . This can be done using the generic algorithm by Cohen et al . ( 2020 ) who employ generative networks to compute barycenters w.r.t . arbitrary discrepancies . They test their method with the maximum mean discrepancy ( MMD ) and Sinkhorn divergence . This approach suffers from the usual limitations of generative models such as mode collapse . Applying it to W2 barycenters requires estimation of W22 ( Pn , P ) . Fan et al . ( 2020 ) test this approach using the minimax method by Makkuva et al . ( 2019 ) , but they end up with a challenging min-max-min problem . • POTENTIAL-BASED OPTIMIZATION : Li et al . ( 2020 ) recover the optimal potentials { ψ∗n } via a non-minimax regularized dual formulation . No generative model is needed : the barycenter is recovered by pushing forward measures using gradients of potentials or by barycentric projection . 4 METHODS . Inspired by Li et al . ( 2020 ) we use a potential-based approach and recover the barycenter by using gradients of the potentials as pushforward maps . The main differences are : ( 1 ) we restrict the potentials to be convex , ( 2 ) we enforce congruence via a regularizing term , and ( 3 ) our formulation does not introduce bias , meaning the optimal solution of our formulation gives the true barycenter . 4.1 DERIVING THE DUAL PROBLEM . Let P be the true barycenter . Our goal is to recover the optimal potentials { ψ∗n , ψ∗n } mapping the input measures Pn into P. To start , we express the barycenter objective ( 4 ) after substituting the dual formulation ( 2 ) : N∑ n=1 αnW22 ( Pn , P ) = [ N∑ n=1 αn ∫ RD ‖x‖2 2 dPn ( x ) ] + ∫ RD ‖y‖2 2 dP ( y ) − min { ψn } ∈Conv [ N∑ n=1 αn ∫ RD ψn ( x ) dPn ( x ) + N∑ n=1 αn ∫ RD ψn ( y ) dP ( y ) ] ( 6 ) The minimum is attained not just among convex potentials { ψn } , but among congruent potentials ( see discussion under ( 5 ) ) ; thus , we can add the constraint that { ψn } are congruent to ( 6 ) . Hence , N∑ n=1 αnW22 ( Pn , P ) = [ N∑ n=1 αn ∫ RD ‖x‖2 2 dPn ( x ) ] − min { ψn } congruent [ N∑ n=1 αn ∫ RD ψn ( y ) dPn ( y ) ︸ ︷︷ ︸ MultiCorr ( { αn , Pn } | { ψn } ) ] . ( 7 ) To transition from ( 6 ) to ( 7 ) , we used the fact that for congruent { ψn } we have∑N n=1 αnψn ( x ) = ‖x‖2 2 , so ∑N n=1 ∫ RD αnψn ( y ) dP ( y ) = ∫ RD ‖y‖2 2 dP ( y ) . We call the value inside the minimum in ( 7 ) the multiple correlation of { Pn } with weights { αn } w.r.t . potentials { ψn } . Notice that the true barycenter P appears nowhere on the right side of ( 7 ) . Thus the optimal potentials { ψ∗n } can be recovered by solving the following min { ψn } congruent MultiCorr ( { αn , Pn } | { ψn } ) = min { ψn } congruent [ N∑ n=1 αn ∫ RD ψn ( y ) dPn ( y ) ] . ( 8 )

Summary: The paper derives the barycenter mapping problem as an optimization over *congruent* convex functions---each convex potential corresponding to a component distribution.  Congruency is a property on the set of optimal potential functions that ties them together.  However, this optimization is quite challenging and so the paper derives an principled objective function that includes two regularization terms.  The first regularization term encourages congruency of the set of convex functions and can be seen as a variational bound on an ideal congruency regularization.  The second regularization term encourages the pairs of convex functions to be conjugate.  The paper proves that the optimal solution of this objective is the true potentials and thus no bias is introduced.  The proposed approach is demonstrated on the tasks of generative modeling (2-256 dimensions), posterior inference, and color pallete barycenters (3D)