Learning Embeddings into Entropic Wasserstein Spaces

1 INTRODUCTION . Learned embeddings form the basis for many state-of-the-art learning systems . Word embeddings like word2vec ( Mikolov et al. , 2013 ) , GloVe ( Pennington et al. , 2014 ) , fastText ( Bojanowski et al. , 2017 ) , and ELMo ( Peters et al. , 2018 ) are ubiquitous in natural language processing , where they are used for tasks like machine translation ( Neubig et al. , 2018 ) , while graph embeddings ( Nickel et al. , 2016 ) like node2vec ( Grover & Leskovec , 2016 ) are used to represent knowledge graphs and pre-trained image models ( Simon et al. , 2016 ) appear in many computer vision pipelines . An effective embedding should capture the semantic structure of the data with high fidelity , in a way that is amenable to downstream tasks . This makes the choice of a target space for the embedding important , since different spaces can represent different types of semantic structure . The most common choice is to embed data into Euclidean space , where distances and angles between vectors encode their levels of association ( Mikolov et al. , 2013 ; Weston et al. , 2011 ; Kiros et al. , 2014 ; Mirzazadeh et al. , 2014 ) . Euclidean spaces , however , are limited in their ability to represent complex relationships between inputs , since they make restrictive assumptions about neighborhood sizes and connectivity . This drawback has been documented recently for tree-structured data , for example , where spaces of negative curvature are required due to exponential scaling of neighborhood sizes ( Nickel & Kiela , 2017 ; 2018 ) . In this paper , we embed input data as probability distributions in a Wasserstein space . Wasserstein spaces endow probability distributions with an optimal transport metric , which measures the distance traveled in transporting the mass in one distribution to match another . Recent theory has shown that Wasserstein spaces are quite flexible—more so than Euclidean spaces—allowing a variety of other metric spaces to be embedded within them while preserving their original distance metrics . As such , they make attractive targets for embeddings in machine learning , where this flexibility might capture complex relationships between objects when other embeddings fail to do so . Unlike prior work on Wasserstein embeddings , which has focused on embedding into Gaussian distributions ( Muzellec & Cuturi , 2018 ; Zhu et al. , 2018 ) , we embed input data as discrete distributions supported at a fixed number of points . In doing so , we attempt to access the full flexibility of Wasserstein spaces to represent a wide variety of structures . Optimal transport metrics and their gradients are costly to compute , requiring the solution of a linear program . For efficiency , we use an approximation to the Wasserstein distance called the Sinkhorn divergence ( Cuturi , 2013 ) , in which the underlying transport problem is regularized to make it more tractable . While less well-characterized theoretically with respect to embedding capacity , the Sinkhorn divergence is computed efficiently by a fixed-point iteration . Moreover , recent work has shown that it is suitable for gradient-based optimization via automatic differentiation ( Genevay et al. , 2018b ) . To our knowledge , our work is the first to explore embedding properties of the Sinkhorn divergence . We empirically investigate two settings for Wasserstein embeddings . First , we demonstrate their representational capacity by embedding a variety of complex networks , for which Wasserstein embeddings achieve higher fidelity than both Euclidean and hyperbolic embeddings . Second , we compute Wasserstein word embeddings , which show retrieval performance comparable to existing methods . One major benefit of our embedding is that the distributions can be visualized directly , unlike most embeddings , which require a dimensionality reduction step such as t-SNE before visualization . We demonstrate the power of this approach by visualizing the learned word embeddings . 2 PRELIMINARIES . 2.1 OPTIMAL TRANSPORT AND WASSERSTEIN DISTANCE . The p-Wasserstein distance between probability distributions µ and ν over a metric space X is Wp ( µ , ν ) = ( inf π∈Π ( µ , ν ) ∫ X×X d ( x1 , x2 ) p dπ ( x1 , x2 ) ) 1 p , ( 1 ) where the infimum is taken over transport plans π that distribute the mass in µ to match that in ν , with the p-th power of the ground metric d ( x1 , x2 ) on X giving the cost of moving a unit of mass from support point x1 ∈ X underlying distribution µ to point x2 ∈ X underlying ν . The Wasserstein distance is the cost of the optimal transport plan matching µ and ν ( Villani , 2003 ) . In this paper , we are concerned with discrete distributions supported on finite sets of points in Rn : µ = M∑ i=1 uiδx ( i ) and ν = N∑ i=1 viδy ( i ) . ( 2 ) Here , u and v are vectors of nonnegative weights summing to 1 , and { x ( i ) } Mi=1 , { y ( i ) } Ni=1 ⊂ Rn are the support points . In this case , the transport plan π matching µ and ν in Equation 1 becomes discrete as well , supported on the product of the two support sets . Define D ∈ RM×N+ to be the matrix of pairwise ground metric distances , with Dij = d ( x ( i ) , y ( j ) ) . Then , for discrete distributions , Equation 1 is equivalent to solving the following : Wp ( µ , ν ) p = min T≥0 tr ( DpT > ) subject to T1 = u , T > 1 = v , ( 3 ) with Tij giving the transported mass between xi and yj . The power Dp is taken elementwise . 2.2 SINKHORN DIVERGENCE . Equation 3 is a linear program that can be challenging to solve in practice . To improve efficiency , recent learning algorithms use an entropic regularizer proposed by Cuturi ( 2013 ) . The resulting Sinkhorn divergence solves a modified version of Equation 3 : Wλp ( µ , ν ) p = min T≥0 tr ( DpT > ) + λtr ( T ( log ( T ) − 11 > ) > ) s.t . T1 = u , T > 1 = v , ( 4 ) where log ( · ) is applied elementwise and λ ≥ 0 is the regularization parameter . For λ > 0 , the optimal solution takes the form T ∗ = ∆ ( r ) exp ( −Dp/λ ) ∆ ( c ) , where ∆ ( r ) and ∆ ( c ) are diagonal matrices with diagonal elements r and c , resp . Rather than optimizing over matrices T , one can optimize for r and c , reducing the size of the problem to M +N . This can be solved via matrix balancing , starting from an initial matrix K : = exp ( −D p λ ) and alternately projecting onto the marginal constraints until convergence : r ← u./Kc c← v./K > r . ( 5 ) Here , ./ denotes elementwise division for vectors . Beyond simplicity of implementation , Equation 5 has an additional advantage for machine learning : The steps of this algorithm are differentiable . With this observation in mind , Genevay et al . ( 2018b ) incorporate entropic transport into learning pipelines by applying automatic differentiation ( back propagation ) to a fixed number of Sinkhorn iterations . 2.3 WHAT CAN WE EMBED IN THEORY ? . Given two metric spacesA and B , an embedding ofA into B is a map φ : A → B that approximately preserves distances , in the sense that the distortion is small : LdA ( u , v ) ≤ dB ( φ ( u ) , φ ( v ) ) ≤ CLdA ( u , v ) , ∀u , v ∈ A , ( 6 ) for some uniform constants L > 0 and C ≥ 1 . The distortion of the embedding φ is the smallest C such that Equation 6 holds . One can characterize how “ large ” a space is ( its representational capacity ) by the spaces that embed into it with low distortion . In practical terms , this capacity determines the types of data ( and relationships between them ) that can be well-represented in the embedding space . Rn with the Euclidean metric , for example , embeds into the L1 metric with low distortion , while the reverse is not true ( Deza & Laurent , 2009 ) . We do not expect Manhattan-structured data to be well-represented in Euclidean space , no matter how clever the mapping . Wasserstein spaces are very large : Many spaces can embed into Wasserstein spaces with low distortion , even when the converse is not true . Wp ( A ) , for A an arbitrary metric space , embeds any product space An , for example ( Kloeckner , 2010 ) , via discrete distributions supported at n points . Even more generally , certain Wasserstein spaces are universal , in the sense that they can embed arbitrary metrics on finite spaces . W1 ( ` 1 ) is one such space ( Bourgain , 1986 ) , and it is still an open problem to determine if W1 ( Rk ) is universal for any k < +∞ . Recently it has been shown that every finite metric space embeds the 1p power of its metric into Wp ( R 3 ) , p > 1 , with vanishing distortion ( Andoni et al. , 2015 ) . A hopeful interpretation suggests thatW1 ( R3 ) may be a plausible target space for arbitrary metrics on symbolic data , with a finite set of symbols ; we are unaware of similar universality results for Lp or hyperbolic spaces , for example . The reverse direction , embedding Wasserstein spaces into others , is well-studied in the case of discrete distributions . Theoretical results in this domain are motivated by interest in efficient algorithms for approximating Wasserstein distances by embedding into spaces with easily-computed metrics . In this direction , low-distortion embeddings are difficult to find . W2 ( R3 ) , for example , is known not to embed into L1 ( Andoni et al. , 2016 ) . Some positive results exist , nevertheless . For a Euclidean ground metric , for example , the 1-Wasserstein distance can be approximated in a wavelet domain ( Shirdhonkar & Jacobs , 2008 ) or by high-dimensional embedding into L1 ( Indyk & Thaper , 2003 ) . In §4 , we empirically investigate the embedding capacity of Wasserstein spaces , by attempting to learn low-distortion embeddings for a variety of input spaces . For efficiency , we replace the Wasserstein distance by its entropically-regularized counterpart , the Sinkhorn divergence ( §2.2 ) . The embedding capacity of Sinkhorn divergences is previously unstudied , to our knowledge , except in the weak sense that the approximation error with respect to the Wasserstein distance vanishes with the regularizer taken to zero ( Carlier et al. , 2017 ; Genevay et al. , 2018a ) . 2.4 RELATED WORK . While learned vector space embeddings have a long history , there is a recent trend in the representation learning community towards more complex target spaces , such as spaces of probability distributions ( Vilnis & McCallum , 2015 ; Athiwaratkun & Wilson , 2018 ) , Euclidean norm balls ( Mirzazadeh et al. , 2015 ; Mirzazadeh , 2017 ) , Poincaré balls ( Nickel & Kiela , 2017 ) , and Riemannian manifolds ( Nickel & Kiela , 2018 ) . From a modeling perspective , these more complex structures assist in representing uncertainty about the objects being embedded ( Vilnis & McCallum , 2015 ; Bojchevski & Günnemann , 2018 ) as well as complex relations such as inclusion , exclusion , hierarchy , and ordering ( Mirzazadeh et al. , 2015 ; Vendrov et al. , 2015 ; Athiwaratkun & Wilson , 2018 ) . In the same vein , our work takes probability distributions in a Wasserstein space as our embedding targets . The distance or discrepancy measure between target structures is a major defining factor for a representation learning model . Lp distances as well as angle-based discrepancies are fairly common ( Mikolov et al. , 2013 ) , as is the KL divergence ( Kullback & Leibler , 1951 ) , when embedding into probability distributions . For distributions , however , the KL divergence and Lp distances are problematic , in the sense that they ignore the geometry of the domain of the distributions being compared . For distributions with disjoint support , for example , these divergences do not depend on the separation between the supports . Optimal transport distances ( Villani , 2008 ; Cuturi , 2013 ; Peyré & Cuturi , 2017 ; Solomon , 2018 ) , on the other hand , explicitly account for the geometry of the domain . Hence , models based on optimal transport are gaining popularity in machine learning ; see ( Rubner et al. , 1998 ; Courty et al. , 2014 ; Frogner et al. , 2015 ; Kusner et al. , 2015 ; Arjovsky et al. , 2017 ; Genevay et al. , 2018b ; Claici et al. , 2018 ; Singh et al. , 2018 ) for some examples . Learned embeddings into Wasserstein spaces are relatively unexplored . Recent research proposes embedding into Gaussian distributions ( Muzellec & Cuturi , 2018 ; Zhu et al. , 2018 ) . Restricting to parametric distributions enables closed-form expressions for transport distances , but the resulting representation space may lose expressiveness . We note that Courty et al . ( 2018 ) study embedding in the opposite direction , from Wasserstein into Euclidean space . In contrast , we learn to embed into the space of discrete probability distributions endowed with the Wasserstein distance . Discrete distributions are dense inW2 ( Kloeckner , 2012 ; Brancolini et al. , 2009 ) .

Summary: The paper proposes embedding the data into low-dimensional Wasserstein spaces. These spaces are larger and more flexible than Euclidean spaces and thus, can capture the underlying structure of the data more accurately. However, the paper simply uses the automatic differentiation to calculate the gradients. Thus, it offers almost no theoretical contribution (for instance, how to calculate these embeddings more efficiently (e.g. faster or more efficient calculation of the Sinkhorn divergence), how to motivate loss functions than can benefit the structure of the Wasserstein spaces, what the interpretation of phi(x) is for each problem, e. g. word embedding, etc.). Additionally, the experiments are unclear and need further improvement. For instance, which method is used to find the Euclidean embedding of the datasets? Have you tried any alternative loss functions for the discussed problems? Are these embeddings useful (classification accuracy, other distortion measure)? How does the 2-D visualizations compare to DR methods such as t-SNE? What is complexity of the method? How does the runtime compare to similar methods?