Generalizing Point Embeddings using the Wasserstein Space of Elliptical Distributions
Embedding complex objects as vectors in low dimensional spaces is a longstanding problem in machine learning. We propose in this work an extension of that approach, which consists in embedding objects as elliptical probability distributions, namely distributions whose densities have elliptical level sets. We endow these measures with the 2-Wasserstein metric, with two important benefits: (i) For such measures, the squared 2-Wasserstein metric has a closed form, equal to a weighted sum of the squared Euclidean distance between means and the squared Bures metric between covariance matrices. The latter is a Riemannian metric between positive semi-definite matrices, which turns out to be Euclidean on a suitable factor representation of such matrices, which is valid on the entire geodesic between these matrices. (ii) The 2-Wasserstein distance boils down to the usual Euclidean metric when comparing Diracs, and therefore provides a natural framework to extend point embeddings. We show that for these reasons Wasserstein elliptical embeddings are more intuitive and yield tools that are better behaved numerically than the alternative choice of Gaussian embeddings with the Kullback-Leibler divergence. In particular, and unlike previous work based on the KL geometry, we learn elliptical distributions that are not necessarily diagonal. We demonstrate the advantages of elliptical embeddings by using them for visualization, to compute embeddings of words, and to reflect entailment or hypernymy.
1 Introduction must comply with sophisticated desiderata. Naturally, more general and flexible approaches in which the embedding space Y needs not be Euclidean can be considered, for instance in generalized MDS on the sphere [Maron et al., 2010], on surfaces [Bronstein et al., 2006], in spaces of trees [Bȃdoiu et al., 2007, Fakcharoenphol et al., 2003 or, more recently, computed in the Poincaré hyperbolic space [Nickel and Kiela, 2017].
Probabilistic Embeddings. Our work belongs to a recent trend, pioneered by Vilnis and McCallum, who proposed to embed data points as probability measures in R d [2015], and therefore generalize point embeddings. Indeed, point embeddings can be regarded as a very particular-and degeneratecase of probabilistic embedding, in which the uncertainty is infinitely concentrated on a single point (a Dirac). Probability measures can be more spread-out, or event multimodal, and provide therefore an opportunity for additional flexibility. Naturally, such an opportunity can only be exploited by defining a metric, divergence or dot-product on the space (or a subspace thereof) of probability measures. Vilnis and McCallum proposed to embed words as Gaussians endowed either with the KullbackLeibler (KL) divergence or the expected likelihood kernel [Jebara et al., 2004]. The Kullback-Leibler and expected likelihood kernel on measures have, however, an important drawback: these geometries do not coincide with the usual Euclidean metric between point embeddings when the variances of these Gaussians collapse. Indeed, the KL divergence and the`2 distance between two Gaussians diverges to 1 or saturates when the variances of these Gaussians become small. To avoid numerical instabilities arising from this degeneracy, Vilnis and McCallum must restrict their work to diagonal covariance matrices. In a concurrent approach, Singh et al. represent words as distributions over their contexts in the optimal transport geometry [Singh et al., 2018].
Contributions. We propose in this work a new framework for probabilistic embeddings, in which point embeddings are seamlessly handled as a particular case. We consider arbitrary families of elliptical distributions, which subsume Gaussians, and also include uniform elliptical distributions, which are arguably easier to visualize because of their compact support. Our approach uses the 2-Wasserstein distance to compare elliptical distributions. The latter can handle degenerate measures, and both its value and its gradients admit closed forms [Gelbrich, 1990], either in their natural Riemannian formulation, as well as in a more amenable local Euclidean parameterization. We provide numerical tools to carry out the computation of elliptical embeddings in different scenarios, both to optimize them with respect to metric requirements (as is done in multidimensional scaling) or with respect to dot-products (as shown in our applications to word embeddings for entailment, similarity and hypernymy tasks) for which we introduce a proxy using a polarization identity.
The paper proposes using elliptical distribution (specifically uniform elliptical distributions) as probability measures for representing data instances in an embedding space. The work is inspired by prior work where the KL divergence was used for Gaussian distribution based embeddings. A motivating factor versus previous work is ensuring consistent behavior as the distributions collapse to Dirac delta functions. The papers use the 2-Wasserstein divergence/distance (motivated by its connection to optimal transport) as the metric for the embedding space. This divergence/distance has a known closed form expression for uniform elliptical distributions that relies on the Bures/Kakutani distance metric. The paper then explores manifold optimizations of the divergence in terms of the covariance matrix, and uses a factored parameterization of the covariance matrix for the optimization. To compute the gradient of the Bures distance a Newton-Schulz algorithm is used to find the matrix square roots and their inverses, which is a more parallelizable manner than the eigendecompostion. The approach is applied to a simple multidimensional scaling example and then to word embedding data sets that have previously been used with other probability distribution based embeddings.   Strength: A novel idea that is fairly well conveyed. The discussion of the manifold optimization and Figure 2 are useful. The optimization of the gradient implementation seems to be well designed. It seems like a fairly complete piece of work with the potential for further refinement and application.  Weaknesses: The text should include more direct motivation on why the added flexibility of an elliptical embedding is useful for modeling versus standard point embedding, especially relating to the hypernymy experiment.    The equation between lines 189 and 190 appears to be wrong. By Equation 2, the trace term should be divided by $d+2$, unless the covariances are pre-scaled.   The arguments (including Figure 3) for using the precision versus the covariance for visualization could be removed. The counterintuitive nature appears to be subjective and distracts from the modeling purpose of a probability distribution based embedding.  Reference for the tractable (computable) form of Bures distance should be Uhlmann. A. Uhlmann, "The “transition probability” in the state space of a∗-algebra," Rep. Math. Phys. 9, 273 (1976)  Some issues with typesetting appear to be problems with inline figures and sectioning.  Minor errors and corrections: Line 25 "are extremely small (2 or 3 numbers" -> "can be represented in two or three dimensional space.  Line 37 add comma around the phrase "or more ... $\langle y_i,y_j\rangle$" Line 55 The phrase "test time" is not clear in this context. Line 55 I cannot find any relevant comment in Table 1. Line 63 "scenarii" Line 69 Should notation for $M$ in pseudo-inverse match other cases (i.e., boldface)? Line 117 Typesetting and line number are messed up by inline figure 1. Line 130 $B$ should be $b$, and $C$ should be $B$.  Line 132 Reasoning for the counterintuitive visualizations is that for uniform distributions the divergence decreases if the scale decreases since the density is more concentrated. Line 146 "adressed" Line 152 Here and elsewhere the references appears sometimes with both parentheses and brackets. Line 155 "on the contrary" to what? This seems to be incorrectly copied from the introductory paragraph in Section 3.  Line 157 "to expensive" Line 166 "w.r.t" Line 177 Should be supplementary material or supplement.  Line 183.1 Should be "a general loss function" since $\matchal{L}$ is indicated. Line 214 "matrics" Line 215 "would require to perform" "would require one to perform" Line 233 Elliptical should be lower case Line 236 "is to use their precision matrix their covariances" ? Line 242 $y_i,y_j$ -> $y_i-y_j$ Line 243 Missing comma before $\mathbf{a}_n$ Figure 5 caption fix the initial quotation mark before Bach. Line 289 "asymetry" Line 310 "embedd"  Line 315 Should clarify if this is the Bures dot product or the Bures cosine. Line 323 "interesections"  The reference list is poorly done: lack of capitalization of proper nouns, inconsistent referencing to journal names (some with abbreviations some without, some missing title case), some URL s and ISSN inserted sporadically, and Inconsistent referencing to articles on arXiv.   ------ Author rebuttal response ----- The author rebuttal helped clarify and I have raised my score and edited the above.   The authors have promised to clarify the dot product equation that appears inconsistent to me.   I agree with Reviewer 3 that the visualization issue is not convincing and would caution elaborating on this somewhat tangential issue more. The original intuition was that the distance between a point and a uniform ellipsoid increase as the scale of the latter increases. However, with the additive nature of the similarity metric (dot product) used for the word embedding this doesn't appear to be the case. Switching the covariance with precision for the uniform elliptical embeddings seems to make it more confusing for Figure 5, where the embedding is trained using the dot product rather than distance.  I also hope the authors take care in carefully proofreading the manuscript and adhering to NIPS template for typesetting.