The Unbalanced Gromov Wasserstein Distance: Conic Formulation and Relaxation

1 Introduction . Comparing data distributions on different metric spaces is a basic problem in machine learning . This class of problems is for instance at the heart of surfaces [ Bronstein et al. , 2006 ] or graph matching [ Xu et al. , 2019 ] ( equipping the surface or graph with its associated geodesic distance ) , regression problems in quantum chemistry [ Gilmer et al. , 2017 ] ( viewing the molecules as distributions of points in R3 ) and natural language processing [ Grave et al. , 2019 , Alvarez-Melis and Jaakkola , 2018 ] ( where texts in different languages are embedded as points distributions in different vector spaces ) . Metric measure spaces . The mathematical way to formalize these problems is to model the data as metric measure spaces ( mm-spaces ) . A mm-space is denoted as X = ( X , d , µ ) where X is a complete separable set endowed with a distance d and a positive Borel measure µ ∈M+ ( X ) . For instance , if X = ( xi ) i is a finite set of points , then µ = ∑ imiδxi ( here δxi is the Dirac mass at xi ) is simply a set of positive weights mi = µ ( { xi } ) ≥ 0 associated to each point xi , which accounts for its mass or importance . For instance , setting some mi to 0 is equivalent to removing the point xi . We 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . refer to Sturm [ 2012 ] for a mathematical account on the theory of mm-spaces . In all the applications highlighted above , it makes sense to perform the comparisons up to isometric transformations of the data . Two mm-spaces X = ( X , dX , µ ) and Y = ( Y , dY , ν ) are considered to be equal ( denoted X ∼ Y ) if they are isometric , meaning that there is a bijection ψ : spt ( µ ) → spt ( ν ) ( where spt ( µ ) is the support of µ ) such that dX ( x , y ) = dY ( ψ ( x ) , ψ ( y ) ) and ψ ] µ = ν . Here ψ ] is the push-forward operator , so that ψ ] µ = ν is equivalent to imposing ν ( A ) = µ ( ψ−1 ( A ) ) for any set A ⊂ Y . For discrete spaces where µ = ∑ imiδxi , then one should have ν = ψ ] µ = ∑ imiδψ ( xi ) . As highlighted by Mémoli [ 2011 ] , considering mm-spaces up to isometry is a powerful way to formalize and analyze a wide variety of problems such as matching , regression and classification of distributions of points belonging to different spaces . Most often , the objects of interest come with a natural distance such as an intrinsic or extrinsic distance and the uniform measure is the usual choice to make mm-spaces widely applicable . The key to unlock all these problems is the computation of a distance between mm-spaces up to isometry . So far , existing distances ( reviewed below ) assume that µ is a probability distribution , i.e . µ ( X ) = 1 . This constraint is not natural and sometimes problematic for most of the practical applications to machine learning . The goal of this paper is to alleviate this restriction . We define for the first time a class of distances between unbalanced metric measure spaces , these distances being upper-bounded by divergences which can be approximated by an efficient numerical scheme . Csiszár divergences The simplest case is when X = Y and one simply ignores the underlying metric . One can then use Csiszár divergences ( or ϕ-divergences ) , which perform a pointwise comparison ( in contrast with optimal transport distances , which perform a displacement comparison ) . It is defined using an entropy function ϕ : R+ → [ 0 , +∞ ] , which is a convex , lower semi-continuous , positive function with ϕ ( 1 ) = 0 . The Csiszár ϕ-divergence reads Dϕ ( µ|ν ) , ∫ X ϕ ( dµ dν ) dν + ϕ′∞ ∫ X dµ⊥ , where µ = dµdν ν + µ ⊥ is called the Radon-Nikodym or the Lebesgue decomposition of µ with respect to ν and ϕ′∞ = limr→∞ ϕ ( r ) /r ∈ R ∪ { +∞ } is called the recession constant . This divergence Dϕ is convex , positive , 1-homogeneous and weak * lower-semicontinuous , see Liero et al . [ 2015 ] for details . Particular instances of ϕ-divergences are Kullback-Leibler ( KL ) for ϕ ( r ) = r log ( r ) − r + 1 ( note that ϕ′∞ =∞ ) and Total Variation ( TV ) for ϕ ( r ) = |r − 1| . Balanced and unbalanced optimal transport . If the common embedding space X is equipped with a distance d ( x , y ) , one can use more elaborated methods such as optimal transport ( OT ) distances , which are computed by solving convex optimization problems . This type of methods has proven useful for ML problems as diverse as domain adaptation [ Courty et al. , 2014 ] , supervised learning over histograms [ Frogner et al. , 2015 ] and unsupervised learning of generative models [ Arjovsky et al. , 2017 ] . In this case , the extension from probability distributions to arbitrary positive measures ( µ , ν ) ∈ M+ ( X ) 2 is now well understood and corresponds to the theory of unbalanced OT . Following Liero et al . [ 2015 ] , Chizat et al . [ 2018a ] , a family of unbalanced Wasserstein distances is defined by solving UW ( µ , ν ) q , inf π∈M ( X×X ) ∫ λ ( d ( x , y ) ) dπ ( x , y ) + Dϕ ( π1|µ ) + Dϕ ( π2|µ ) . ( 1 ) Here ( π1 , π2 ) are the two marginals of the joint distribution π , defined by π1 ( A ) = π ( A× Y ) for A ⊂ X . The mapping λ : R+ → R and exponent q ≥ 1 should be chosen wisely to ensure for instance that UW defines a distance ( see Section 2.2.1 ) . It is frequent to take ρDϕ instead of Dϕ ( i.e . take ψ = ρϕ ) to adjust the strength of the marginals ’ penalization . Balanced OT is retrieved with the convex indicator ϕ = ι { 1 } ( i.e . ϕ ( 1 ) = 0 and ϕ ( x ) = +∞ otherwise ) or by taking the limit ρ → +∞ , which enforces π1 = µ and π2 = ν . When 0 < ρ < +∞ , unbalanced OT operates a trade-off between transportation and creation of mass , which is crucial to be robust to outliers in the data and to cope with mass variations in the modes of the distributions . For supervised tasks , the value of ρ should be cross-validated to obtain the best performances . Its use is gaining popularity in applications , such as medical imaging registration [ Feydy et al. , 2019a ] , videos [ Lee et al. , 2019 ] , generative learning [ Balaji et al. , 2020 ] and gradient flow to train neural networks [ Chizat and Bach , 2018 , Rotskoff et al. , 2019 ] . Furthermore , existing efficient algorithms for balanced OT extend to this unbalanced problem . In particular Sinkhorn ’ s iterations , introduced in ML for balanced OT by Cuturi [ 2013 ] , extend to unbalanced OT [ Chizat et al. , 2018b , Séjourné et al. , 2019 ] , as detailed in Section 3 . The Gromov-Wasserstein distance and its applications . The Gromov-Wasserstein ( GW ) distance [ Mémoli , 2011 , Sturm , 2012 ] generalizes the notion of OT to the setting of mm-spaces up to isometries . It replaces the linear cost ∫ λ ( d ) dπ of OT by a quadratic function . It reads GW ( X , Y ) q , min π∈M+ ( X×Y ) { ∫ λ ( |dX ( x , x′ ) − dY ( y , y′ ) | ) dπ ( x , y ) dπ ( x′ , y′ ) : π1 = µ π2 = ν } . ( 2 ) It is proved in Mémoli [ 2011 ] , Sturm [ 2012 ] that GW defines with λ ( t ) = tq a distance up to isometries on balanced mm-spaces ( i.e . the measures are probability distributions ) . The GW distance is applied successfully in natural language processing for unsupervised translation learning [ Grave et al. , 2019 , Alvarez-Melis and Jaakkola , 2018 ] , in generative learning for objects lying in spaces of different dimensions [ Bunne et al. , 2019 ] and to build VAE for graphs [ Xu et al. , 2020 ] . It has been adapted for domain adaptation over different spaces [ Redko et al. , 2020 ] . It is also a relevant distance to compute barycenters between graphs or shapes [ Vayer et al. , 2018 , Chowdhury and Needham , 2020 ] . When ( X , Y ) are Euclidean spaces , this distance compares distributions up to rigid isometry , and is closely related ( but not equal ) to metrics defined by procrustes analysis [ Grave et al. , 2019 , Alvarez-Melis et al. , 2019 ] . The problem ( 2 ) is non convex because the quadratic form∫ λ ( |dX − dY | ) dπ ⊗ π is not positive in general . It is in fact closely related to quadratic assignment problems [ Burkard et al. , 1998 ] , which are used for graph matching problems , and are known to be NP-hard in general . Nevertheless , non-convex optimization methods have been shown to be successful in practice to use GW distances for ML problems . This includes for instance alternating minimization [ Mémoli , 2011 , Redko et al. , 2020 ] and entropic regularization [ Peyré et al. , 2016 , Gold and Rangarajan , 1996 ] . Related works and contributions . The concomitant work of De Ponti and Mondino [ 2020 ] extends the Lp transportation distance defined in Sturm et al . [ 2006 ] to unbalanced mm-spaces and studies its geometric properties . This distortion distance is not equivalent to the GW distance , and is more difficult to estimate numerically because it explicitly imposes a triangle inequality constraint in the optimization problem . The work of Chapel et al . [ 2020 ] relaxes the GW distance to the unbalanced setting by hybridizing GW with partial OT [ Figalli , 2010 ] for unsupervised labeling . It ressembles one particular setting of our formulation , but with some important differences , detailed in Section 2 . Our construction is also connected to partial matching methods , which find numerous applications in graphics and vision [ Cosmo et al. , 2016 ] . In particular , Rodola et al . [ 2012 ] introduces a mass conservation relaxation of the GW problem . The two main contributions of this paper are the definition of two formulations relaxing the GW distance . The first one is called the Unbalanced Gromov-Wasserstein ( UGW ) divergence and can be computed efficiently on GPUs . The second one is called the Conic Gromov-Wasserstein distance ( CGW ) . It is proved to be a distance between mm-spaces endowed with positive measures up to isometries , as stated in Theorem 1 which is the main theoretical result of this paper . We also prove in Theorem 1 that UGW can be used as a surrogate upper-bounding CGW . We present those concepts and their properties in Section 2 . We also detail in Section 3 an efficient computational scheme for a particular setting of UGW . This method computes an approximate stationary point of a biconvex relaxation of our formulations . Even though it is a lower bound of the original problem , we provide in Theorem 3 conditions ensuring the tightness of this relaxation in many cases of interest . The algorithm leverages the strength of entropic regularization and the Sinkhorn algorithm , namely that it is GPU-friendly and defines smooth loss functions amenable to back-propagation for ML applications . Section 4 provides some numerical experiments to highlight the qualitative behavior of this algorithm and its ability to cope with outliers and mass variations in the modes of the distributions . We illustrate numerically the tightness of the relation between UGW and CGW , showing that UGW is a reasonnable proxy of a distance , at least locally . We provide an application of our divergence in the positive unlabeled learning setting , using domain adaptation data , and display results which are at par or outperform the computable competitor Chapel et al . [ 2020 ] .

Summary: This paper proposes two unbalanced Gromov-Wasserstein formulations, namely Unbalanced Gromov-Wasserstein (UGW) divergence and Conic Gromov-Wasserstein (CGW) distance. Both of them are positive and definite. UGW has a scalable and GPU-friendly algorithm, which is applicable to large learning problems. CGW is a distance between mm-spaces up to isometry. It is theoretically and empirically shown that UGW can be used as a surrogate upper-bounding CGW. The authors then apply UGW to positive unlabeled learning with Caltech office dataset and achieve comparable or better performance than Chapel et al., 2020.