Manifold clustering in the embedding space using UMAP and GMM
By: Gianmario Spacagna
Tags: Gmm, Umap, Clustering, Dimensionality Reduction, Manifold Learning

Manifold clustering in the embedding space using UMAP and GMM

Photo by DIMA VALENTINA on Unsplash

In the previous article Extracting rich embedding features from pictures using PyTorch and ResNeXt-WSL we have seen how to represent pictures into a multi-dimensional numerical embedding space. We have also seen the effectiveness of the embedding space to represent similar pictures closely to each other. In this tutorial, we will see a few clustering techniques that are suitable for discovering and identifying the manifolds in our dataset.

Moreover, the proposed clustering technique is also motivated by the “document embedding averaging” that will be described in the next article.

Dimensionality reduction via Uniform Manifold Approximation and Projection (UMAP)

Dimensionality reduction is not just used for data visualization, but it is a fundamental step for clustering algorithms due to the “ curse of dimensionality “. In other words, if the number of dimensions increases then most of the points will start to look as similar and as different from each other across at least a few of those dimensions. The effect is that there is no clear structure to follow resulting in a random grouping of the data points.

The unsupervised dimensionality reduction techniques are divided into two families: Linear Projection and Manifold Learning.

The main difference of manifold learning with linear projections (e.g. PCA, SVD) is that it can handle non-linear relationships in the data and it is very effective for clustering groups of similar data points preserving their relative proximities. For our purposes and given the nature of our data (the embedding space generated by a deep convolutional neural network), we do not consider any linear projection technique as suitable.

In the manifold learning family, there are two main competitors: t-SNE and UMAP.

t-SNE 3d projection of COCO detection 2017 pictures colored by supercategory

UMAP 3d projection of COCO detection 2017 pictures colored by supercategory

t-SNE vs UMAP 3d projection of COCO detection 2017 pictures colored by supercategory

Uniform Manifold Approximation and Projection (UMAP) is a general-purpose manifold learning and dimension reduction algorithm.

T-distributed Stochastic Neighbour Embedding (t-SNE) is an algorithm that generates a low-dimensional graph trying to keep similar instances close and dissimilar instances apart.

They sound similar, and in fact from a lot of aspects they are. Nevertheless, let’s summarize a few reasons to prefer UMAP over t-SNE for clustering purposes:

Comparison of t-SNE and UMAP

For a more comprehensive comparison of t-SNE vs. UMAP please refer to the following article: How exactly UMAP works.

For the reasons discussed above, we can conclude that t-SNE is a great visualization tool but UMAP is a more suitable technique for clustering purposes in the case of manifold structures.

Clustering with Gaussian Mixture Model (GMM)

GMM is a probabilistic model that assumes all the data points are generated from a mixture of a finite number of Gaussian distributions with unknown parameters. It can be seen as a generalization of the more popular k-means model. The advantage of using GMM over k-means is that it can represent clusters of different sizes and shapes based on a parametric covariance matrix.

In k-means the clusters are spherical over all of the dimensions and share the same diameter. This is a big limitation when considering a manifold feature space as is the case of transforming a deep convolutional neural network embedding space with UMAP.

Different clustering output using the two models on the iris dataset

(source)

Another distinction is in the interpretation of the clustering output. k-means divide the space into voronoi cells and hard assign each point to the cluster of the closest centroid. GMM, on the other hand, gives us an interpretable output modeling the probability that each data point belong to each cluster. The latter is a desired probability for dealing with fuzzy situations in presence of overlapping clusters or outliers.

GMM parameters

As with k-means, also GMM requires the number of clusters k to be specified.

Moreover, in order for GMM to be able to model arbitrary elliptic shapes in the feature space, the covariance matrix should be “full”. The problem with “full” GMM models is that the degrees of freedom increase quadratically with the dimension of the feature space, risking to overfit the data. There are a few constrained versions of GMM that impose certain properties to the covariance matrix., namely: spherical, diagonal, and tied.

Cluster shapes obtained using different covariance types

(source: scikit-learn)

Spherical is a “diagonal” situation with circular contours (spherical in higher dimensions, whence the name).

is a “diagonal” situation with circular contours (spherical in higher dimensions, whence the name). Diagonal means the contour axes are oriented along the coordinate axes, but otherwise the eccentricities may vary between components.

means the contour axes are oriented along the coordinate axes, but otherwise the eccentricities may vary between components. Full means the components may independently adopt any position and shape.

means the components may independently adopt any position and shape. Tied means they have the same shape, but the shape may be anything.

We are now left with two major parameters to tune: the number of clusters k and the covariance type among the 4 options listed above.

Model selection using BIC and Silhouette scores

Likely, since GMM is a probabilistic model we can calculate the Bayesian Information Criterion (BIC) that is a statistics calculated as the sum of the negative log-likelihood of the model and a penalty term that is a function of the number of data samples and the number of free parameters of the model. The smaller the BIC value the more preferable is the model. Nonetheless, searching for the minimum BIC score may suggest selecting a model with a lot of clusters in front of tiny decreases of the score. That is why a preferred approach is to identify the elbow of the curve that corresponds to the minimum of the second derivative.

The BIC score is comparable among different clustering outputs only if they are representing the same points in the same feature space. That is, we cannot compare data points reduced via PCA with data points reduced via UMAP.

Another technique is the Silhouette score that is an empirical method measuring the consistency of the clusters by comparing how much a point is similar to its cluster (cohesion) compared to the other clusters (separation). The silhouette ranges from −1 to +1, where a high value indicates that the object is well matched to its own cluster and poorly matched to neighboring clusters. If most objects have a high value, then the clustering configuration is appropriate, otherwise the clustering configuration may have too many or too few clusters.

If we plot both the BIC and Silhouette curves as function of the number of clusters k for the 4 different covariance types we obtain the following graph:

GMM model selection using both the BIC and Silhouette scores. Image by Author.

The dot points correspond, respectively, to the elbow and the maximum of the BIC and Silhouette curves.

We can conclude that the ideal number of clusters should be between 30 and 50.

In terms of the covariance type, the tied type minimizes the BIC while there is not strong evidence of worsening results in the Silhouette curve.

The lower BIC score can be explained by the good trade-off between low model complexity and the high likelihood of the points. Moreover, given the nature of the feature space, it does make sense to consider manifolds of irregular but similar shapes.

Tied covariance matrix of the GMM model trained with 40 clusters. Image by Author.

The selected configuration for us will be tied covariance type and 40 clusters.

Visualizing the clusters

In order to visualize the clusters we will re-use the 3D projections that were calculated in Extracting rich embedding features from pictures using PyTorch and ResNeXt-WSL but we will color the points based on the assigned clusters rather than the COCO super categories of the pictures.

First row: PCA 3D projection colored by COCO categories (on the left) or the GMM clusters (on the right). Second row: UMAP 3D projection colored by COCO categories (on the left) or the GMM clusters (on the right).

Comparison with the COCO taxonomy

Predicting COCO annotations via supervised learning

In order to measure the predictive power of the embedding space, let’s try to predict the COCO annotations. We can train a random forest classifier with default parameters on a multi-labels task. Since each picture can have none, one, or many categories annotated, the task consists of predicting whether a given label is present or not in the picture.

I have used 4000 pictures for training and 1000 for the test stratifying on the most frequent label in each picture. The multilabel accuracy ( exact subset of labels match as defined in scikit-learn) would be 16.7% which is not bad considering the high cardinality of the task. If we micro-average all of the label predictions we can perform a binary evaluation:

Confusion matrix for the multi-label random forest classifier

We have achieved 89% precision and 31% recall on all of the possible picture annotations, not bad and not great. We should consider that we only trained on a very small sample of pictures and a few labels had very few occurrences. Nonetheless, the purpose of this tutorial is not to predict COCO categories but rather to show the effectiveness of the embedding features in identifying correct manifolds.

Clusters similarity and consistency

Since we have grouped the COCO pictures into 40 unsupervised clusters, let’s compare our grouping with the categories provided in the COCO taxonomy. We can use the Adjusted Rand Index that is a measure of similarity between two clustering outputs. The higher the score, the higher the consistency between the two groupings.

We obtained the following results:

AdjustedRand(GMM_clusters, COCO_supercategories) = 0.09955

AdjustedRand(GMM_clusters, COCO_categories) = 0.08844

As we could have already observed from the 3D projections, we can conclude that the discovered topicality defined by the manifolds in the data does not match the COCO taxonomy.

Pictures in the cluster

What topic is each cluster representing then?

Let’s print the pictures closest to the centroid in a few sample clusters.

Cluster 0: horses

Cluster 3: dining

Cluster 10: sea and watersports

Cluster 18: bears

Cluster 22: towers

Conclusions

In this tutorial, we have learned how to cluster pictures in their latent embedding space. We first have used UMAP for isolating manifolds and projecting them into a lower-dimensional space. We then used GMM for discovering the high-density areas in the UMAP space. The BIC elbow and Silhouette techniques were used to find the ideal number of clusters as well as the constraints to the covariance matrix. Through the AdjustedRand test, we demonstrated that the data is intrinsically organized into major topics that do not match with the COCO taxonomy. For instance, we found clusters for horses, bears, towers, watersports, people’s dining, and more.

The presented methodology can be used to cluster any dataset that presents high-dimensional manifolds, not just pictures. It is in general suitable for embeddings produced by neural network models.

If instead of using a pre-trained network, you are training your own, you may want to consider a small dimensionality (below 50) such that you may not need any dimensionality reduction before clustering. Other clustering algorithms that work with any kind of shapes, and are not constrained by the Gaussian mixture assumption, are the hierarchical density-based models such as HDBSCAN.

Stay tuned for the next article on how to exploit embedding features and the manifolds clusters to average and represent collections of datapoints (documents) into the same latent space.

You can find the code and the notebooks at https://github.com/gm-spacagna/docem.

If you want to learn more about tuning techniques for data clustering algorithms you can read those articles: Data Clustering? don’t worry about the algorithm and A Distributed Genetic Evolutionary Tuning for Data Clustering: Part 1.