Despite their prevalence, Euclidean embeddings of data are fundamentally limited in their ability to capture latent semantic structures, which need not conform to Euclidean spatial assumptions. Here we consider an alternative, which embeds data as discrete probability distributions in a Wasserstein space, endowed with an optimal transport metric. Wasserstein spaces are much larger and more flexible than Euclidean spaces, in that they can successfully embed a wider variety of metric structures. We propose to exploit this flexibility by learning an embedding that captures the semantic information in the Wasserstein distance between embedded distributions. We examine empirically the representational capacity of such learned Wasserstein embeddings, showing that they can embed a wide variety of complex metric structures with smaller distortion than an equivalent Euclidean embedding. We also investigate an application to word embedding, demonstrating a unique advantage of Wasserstein embeddings: we can directly visualize the high-dimensional embedding, as it is a probability distribution on a low-dimensional space. This obviates the need for dimensionality reduction techniques such as t-SNE for visualization. Learned embeddings form the basis for many state-of-the-art learning systems. Word embeddings like word2vec BID34 , GloVe BID42 , fastText BID5 , and ELMo BID43 are ubiquitous in natural language processing, where they are used for tasks like machine translation BID38 , while graph embeddings BID41 like node2vec BID21 are used to represent knowledge graphs and pre-trained image models BID47 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 BID34 BID56 BID27 BID36 . 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 BID39 BID49 .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 BID37 BID58 , 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 BID15 , 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 BID20 . 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. Several characteristics determine the value and effectiveness of an embedding space for representation learning. The space must be large enough to embed a variety of metrics, while admitting a mathematical description compatible with learning algorithms; additional features, including direct interpretability, make it easier to understand, analyze, and potentially debug the output of a representation learning procedure. Based on their theoretical properties, Wasserstein spaces are strong candidates for representing complex semantic structures, when the capacity of Euclidean space does not suffice. Empirically, entropy-regularized Wasserstein distances are effective for embedding a wide variety of semantic structures, while enabling direct visualization of the embedding.Our work suggests several directions for additional research. Beyond simple extensions like weighting points in the point cloud, one observation is that we can lift nearly any representation space X to distributions over that space W(X ) represented as point clouds; in this paper we focused on the case X = R n . Since X embeds within W(X ) using δ-functions, this might be viewed as a general "lifting" procedure increasing the capacity of a representation. We can also consider other tasks, such as co-embedding of different modalities into the same transport space. Additionally, our empirical results suggest that theoretical study of the embedding capacity of Sinkhorn divergences may be profitable. Finally, following recent work on computing geodesics in Wasserstein space BID45 , it may be interesting to invert the learned mappings and use them for interpolation.
We show that Wasserstein spaces are good targets for embedding data with complex semantic structure. Learns embeddings in a discrete space of probability distributions, using a minimized, regularised version of Wasserstein distances. The paper describes a new embedding method that embeds data to the space of probability measures endowed with the Wasserstein distance.  The paper proposes embedding the data into low-dimensional Wasserstein spaces, which can capture the underlying structure of the data more accurately.