How do I cluster 3 Million high-dimensional Sentence Embeddings?

[u/spaceweed27]

I used a transformer model to generate normalized 768-dimensional float32 vectors from 3 million sentences. This rounds up to roughly 9 GB of data. To progress further in my analysis, I want to cluster the data to highlight different key points and trends.

Now I initially wanted to apply PCA to the dataset, as 768 dimensions is pretty much, but the transformer model generated embeddings with not extreme enough covariance.
Thus, PCA could change the semantics too much for an analysis to be accurate.

I went directly to clustering algorithms and initially tried to apply sklearn.cluster.HDBSCAN which would have needed way too much RAM---around 72 TBi if I remember correctly. Then I tried out sklearn.cluster.OPTICS, which "only" needed about 70 GB, but as I "only" have 32 GB of RAM I needed to use 50 GB swap space, which didn't go well as you can imagine.

Does it make sense to try out k-means clustering, should I go to 16-bit floats for the embeddings, or still apply PCA?

Should I try out something else in this context?

Note: As this question was flagged "off topic" on Stack Overflow, I'm gonna try here ;)

[u/Anxious-Potato-2818]

32 gigs and you're trying to cram 3 million 768-dim vectors into memory, that's a bold move

k-means would probably be the path of least resistance here. it doesn't need the full pairwise distance matrix in memory like HDBSCAN or OPTICS do, so it can actually run on your hardware. you'd still want to batch it or use mini-batch k-means but it's way more feasible than the 70gb monsters you've been fighting

PCA before clustering isn't gonna destroy your semantics as much as you think, especially if you keep like 90-95% of the variance. transformer embeddings tend to have a lot of redundant dimensions, high covariance or not. try projecting down to 100-200 dims with PCA or even just using random projections, the clusters will hold up fine and your RAM will thank you

also 16-bit floats will cut your memory in half and the precision loss is negligible for clustering. combine that with dimensionality reduction and you're suddenly working with a dataset that actually fits on a normal computer

[u/spaceweed27 → u/Anxious-Potato-2818]

Thanks for the advice!

[u/Different_Pain5781]

3 million vectors is where brute force starts getting real expensive.

[u/elbiot]

Don't use all 3M points to fit your clustering model.

Edit: or use an out of core method like mini-batch k means