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: This work introduces a new Wasserstein-2 barycenter computation method. The authors first derive the dual formulation of the Wasserstein-2 barycenter problem, and then parametrize the convex potentials by ICNNs. The congruent and conjugacy conditions are enforced by regularization terms, respectively. They then show that the algorithm can find a good barycenter if the objective function is properly minimized.