Wasserstein t-SNE
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.