Wasserstein t-SNE

TL;DR

Wasserstein t-SNE addresses visualization of hierarchical data where units are distributions over samples by computing pairwise unit distances with the Wasserstein metric and embedding the resulting matrix in 2D using $t$-SNE. The method offers a Gaussian-approximation route for scalable computation, with a tunable $\lambda$ that balances mean and covariance information, and an exact LP-based route for faithful distances at higher cost. Through synthetic HGMM simulations and a real 2017 German election dataset, the approach demonstrates improved clustering and reveals meaningful cross-regional correlation patterns that are not captured by unit means alone. The work provides practical, scalable tools for unit-level visualization in hierarchical datasets, with potential applications across social sciences and biomedical domains, and offers insight into when exact vs. approximate Wasserstein distances are preferable.

Abstract

Scientific datasets often have hierarchical structure: for example, in surveys, individual participants (samples) might be grouped at a higher level (units) such as their geographical region. In these settings, the interest is often in exploring the structure on the unit level rather than on the sample level. Units can be compared based on the distance between their means, however this ignores the within-unit distribution of samples. Here we develop an approach for exploratory analysis of hierarchical datasets using the Wasserstein distance metric that takes into account the shapes of within-unit distributions. We use t-SNE to construct 2D embeddings of the units, based on the matrix of pairwise Wasserstein distances between them. The distance matrix can be efficiently computed by approximating each unit with a Gaussian distribution, but we also provide a scalable method to compute exact Wasserstein distances. We use synthetic data to demonstrate the effectiveness of our Wasserstein t-SNE, and apply it to data from the 2017 German parliamentary election, considering polling stations as samples and voting districts as units. The resulting embedding uncovers meaningful structure in the data.

Figures (10)

• Figure 1: Hierarchical data. Individual samples can be grouped into units. Each unit forms a probability distribution over its samples. In our Wasserstein $t$-SNE approach, units in the dataset are compared using the Wasserstein metric to construct a pairwise distance matrix, which is then embedded in two dimensions using the $t$-SNE algorithm. Units with similar probability distributions end up close together in the 2D embedding.
• Figure 2: Wasserstein distance as a linear program. ( A) The optimal transport map $\gamma$ of two probability distributions $\nu$ (orange) and $\mu$ (blue) is shown. The heatmap represents the cost matrix $C$. ( B) The same distributions can be visualized as a collection of samples, which have different support. The distance between samples $\nu_i$ and $\mu_j$ is given in the cost matrix entry $C_{ij}$. The size of the optimization variable $\gamma$ is then upper bounded by the product of the sample sizes.
• Figure 3: Computation time and accuracy. ( A) Two multivariate Gaussian distributions with $50$ samples each. ( B) The Wasserstein distance between the two probability distributions is computed using a different number of samples. The ground-truth distance is obtained by the closed-form solution and is shown with the dashed black line. The Wasserstein distance estimates using our linear program approach are shown in green (mean and standard deviation over $50$ repetitions). The purple line shows the average runtime.
• Figure 4: Hierarchical Gaussian mixture model (HGMM). This two-dimensional example dataset has $K=4$ classes with $N=100$ units each. The gray points show the samples from all units ($M=15$ samples per unit). Note that some of the units in the red and green classes have similar means, but their covariances are very different.
• Figure 5: Wasserstein $t$-SNE. ( A) This two-dimensional ($F=2$) HGMM was generated using $K=4$ classes with $N=100$ units each ($M=30$ samples per unit). Two pairs of classes have the same distribution of unit means, while two other pairs of classes have the same distribution of unit covariance matrices. ( B) The mean-based embedding ($\lambda=0$) is not able to separate some of the classes. ( C) The Wasserstein embedding ($\lambda=0.5$) separates all four classes. ( D) The covariance-based embedding ($\lambda=1$) is not able to separate some of the classes. ( E) The performance at different values of $\lambda$ was assessed using the kNN accuracy ($k=5$) in the 2D embedding and the adjusted Rand index (ARI) obtained from Leiden clustering of the original distance matrix (kNN graph with $k=5$, resolution parameter $\gamma=0.08$).

Theorems & Definitions (2)

• Definition 1
• Definition 2