Clustering with large number of clusters

Question

I would like to cluster tens of millions of vectors (hidden states of BERT) into something like 20k clusters. Is there a clustering method that can do this in a reasonable time?

Standard K-means algorithm is too memory demanding: requires computing the similarity between all points and all centroids. Batched K-means are also out of the question because the batch size needs to be significantly larger than the number of centroids which again leads to the same memory issues.

I tried hierarchical clustering with hdbscan and fastcluster which should not suffer from this problem, however, it seems to never finish building the hierarchy.

Is there a clustering algorithm that could succeed under these conditions?

Answer 1

The solution I used, in the end, was my implementation of batched K-Means.

Usual implementations of batched K-Means do both the expectation and the maximization step on a single batch. This is not possible in this case becase the data bach must be smaller than the number of clusters.

The solution is to do the expectation step in batches on the entrire dataset, i.e., for each batch compute the cluster assigment given the current centroids and remember the assigmnet. Having the assigment, I can do do the maximization step on the entire dataset.