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: The paper extends the Gromov-Wasserstein (GW) distance framework to two novel relaxations, the unbalanced GW (UGW) and the conic GW (CGW), which both address the question of comparing arbitrary positive measures on different metric spaces. The paper also proves that CGW is upper bounded by UGW, and introduces an algorithm to compute UGW based on Sinkhorn iterations. Numerical experiments are presented to illustrate the performance of the new algorithm.