"People have been doing minibatch k-means for a while, including with an EMA update" : EMA updates are indeed used in clustering, but to our knowledge there is no paper that uses pure EMA updates for k-means clustering on a static dataset. EMA updates are typically used for moving distributions because of their ability to forget the past (as in VQ-VAEs), but their advantage for static datasets had not been empirically demonstrated. If you take a look at appendix E of the paper, we compare EMA to other gradient descent methods and show that EMA is much more stable and has better quality, especially at large batches. Our method also integrates EMA updates in an optimization framework that reliably outperforms Lloyd and previous mini-batch methods.
"the contribution here is that they tried to measure how much better it is": The contribution here is to integrate EMA updates in an optimization framework to provide a method as robust and performant as we could. If you look at previous mini-batch k-means papers, the speedups and the performance is nowhere nearly as good as e-means, which beats Lloyd at 25 epochs in 2 to 3 epochs. In fact, mini-batch k-means is known in the literature as a fast method but with lesser convergence than Lloyd, and we compare to them in Figure 3. The e-means framework allows to be both faster and have better convergence than Lloyd, without the user having to tune the batch or the step.
"How your algorithm scales as more data is present": In the first figure of the paper, we show that e-means scales up to 10 billion points and 16 million clusters.
On communication: workers only exchange per-update statistics (K×d numbers with replicated centroids). You argue that bigger batches have diminishing returns, so that extra local work doesn't cut the number of updates. But that's precisely what our batch law addresses: it returns the largest batch whose loss stays within ~0.2% of the best loss at a fixed number of epochs. In other words, we carefully scale the batch size to avoid that our speedups are driven down by the diminishing returns you mention.
The multi-node regime is in the paper: the 10B run used 32 machines, so synchronization went over the inter-node network, and the bottleneck was still the distance computation, not communication. That's also why an IVF variant, which cuts distance computations, does it in ~23 GPU-hours.
So I'll be honest the reason I said this is that I know there are people who are using it and I just assumed it was common :P I will ask if any of them are relying on published work or what for where they got this from.
> In the first figure of the paper, we show that e-means scales up to 10 billion points and 16 million clusters.
Sorry, I should have clarified. I actually think this is what is saving you: with that many clusters, your arithmetic intensity is inherently really high, so you are truly super compute limited. In this case introducing a bit of comms is cheap relative to the bazillion matmuls you're doing for distance calculation. However if you have fewer centroids then the amount of work you do is much lower between updates. As an extreme example, if you're doing hierarchical k-means you might have as few as 2 centroids, which means you barely do any work, and adding any sort of comms just kills you. Of course you are doing very little math but even if you're just looking at memory bandwidth at that point (because your intensity is like 5 FLOPs/byte or something) it will be done too fast.
(I don't actually know how to solve this, which is why I was super excited to see someone else working on this. I agree that your thing works for more centroids, in fact I know it works because I just assumed it was what you were supposed to do, I just couldn't make it work in this situation.)
I'd be very interested if they have a reference. If you feel like it, you can always send me an e-mail at [email protected] with the references. I'll add them to my benchmarks.
The fewer centroids case is super interesting. Because we found in our experiments that basically if your n/k ratio (n is the dataset size, k the number of clusters) is high, running the clustering as is on your n points is a waste of resource. You should subsample your dataset such that there is about 250 points per cluster and you will get very similar loss at a fraction of the compute. This is why faiss k-means, by default, subsamples datasets to 256 points per cluster if n/k > 256. So I fully agree with your computational remarks, but I think it's a case that should never happen in practice because subsampling your dataset is the solution in those cases.
If you are claiming that your algorithm is "highly parallelizable" and comes with "wall-clock speedups", then the actual metrics that matter here is not number of steps or even basic wall clock time, but how your algorithm scales as more data is present or parallelism is required. This algorithm splits up the minibatch across workers and has them compute local updates, then allreduces it. As the amount of data goes up, the comms become much more expensive (since you have to use a worse transport layer), which means that you're basically locked in a desperate attempt to do more local work before you send your partial update to everyone else. If you attempt this (not explored in the paper) you will find that the improvement in clustering by doing this extra work has diminishing returns–you only reduce the number of updates by a few steps (which is what drives your wall clock time). At scale you end up being heavily comms limited and not super parallel at all.