distributed k-means - clustering datasets that outgrow one machine

One point in this post is eight float64 numbers, which is 64 bytes. A billion of them is 64 GB. That number is the whole stake of the article, and it is arithmetic, not a measurement.

I want to be exact about that up front, because everything below leans on it. There is no billion-point job behind this post. I constructed one on paper: a billion points, eight dimensions each, float64, sitting in object storage. Every quantity I give for that dataset is multiplication that can be redone by hand, and I flag it as arithmetic each time. No stopwatch appears anywhere below.

What the constructed dataset is good for is one question. What changes about k-means when the array will not fit in one machine’s memory? sklearn.cluster.KMeans wants a materialized array. Hand it 64 GB on a 16 GB laptop and the fit call fails inside the allocation, ahead of the first assignment step. The algorithm is unchanged. The machine is the constraint.

If you have called KMeans on a dataframe that fit in RAM, that is the only prerequisite. If you have not, the first section rebuilds the loop out of three points and two centroids, so this reads cold either way. After that I add exactly one new thing per section, and the machinery gets its standard names only once you have watched it work.

One point, one centroid, one distance

Three points on a line: p1 at 1, p2 at 3, p3 at 10. Two centroids on the same line: c1 at 0, c2 at 8. That is the entire dataset. It fits in a sentence, which is why I picked it.

k-means is trying to do one thing with those centroids: place them so that the total squared distance from every point to its own centroid is as small as it can get. That total has a name, inertia, and it is the number the algorithm is grinding down. Hold onto it, because it is the only scoreboard in the room.

The assignment step asks one question per point. Which centroid is nearer?

Take p2 at 3. Its distance to c1 is 3 and its distance to c2 is 5. Three is smaller, so p2 joins c1. Run the same comparison twice more. p1 at 1 sits 1 from c1 and 7 from c2, so it joins c1. p3 at 10 sits 10 from c1 and 2 from c2, so it joins c2.

Now notice what that comparison needed. It needed p2, it needed the two centroids, and it needed nothing further. Not p1. Not p3. Not the other 999,999,997 points in the constructed version. The assignment of a point is a function of the point, the k centroids, and nothing else in the dataset.

That independence is the load-bearing fact of this post. Nothing has been distributed yet, and the fact does not have a name yet either.

The update step then moves each centroid to the mean of the points that chose it. c1 becomes the mean of 1 and 3, which is 2. c2 becomes the mean of 10, which is 10. Both moved, and inertia fell. Repeat the assignment against the new centroids until nothing moves, which for this toy takes one more round.

import numpy as np
from sklearn.cluster import KMeans

X = np.array([[1.0], [3.0], [10.0]])
km = KMeans(n_clusters=2, init=np.array([[0.0], [8.0]]), n_init=1)
km.fit(X)

km.cluster_centers_   # array([[ 2.], [10.]])
km.labels_            # array([0, 0, 1])

Two machines and the same three points

One rung up. Put p1 and p2 on machine A, put p3 on machine B, and leave them where they are.

Send both centroids to both machines. That is two numbers, 0 and 8. Machine A assigns its two points using its two points and those two numbers; machine B does the same for its one. Neither machine needs anything the other is holding, which is the fact from the previous section doing its job.

Then the update, which is where the interesting failure lives. The obvious move is to have each machine compute its own mean and average the results. That is wrong, and the toy shows why in one line. Machine A’s cluster-1 mean is 2, taken over two points. Machine B holds no cluster-1 points, so its cluster-1 mean does not exist. Even in the friendlier case where both machines hold some, averaging the two means weights a machine holding one point the same as a machine holding a million.

The mean is not composable, but the pair of a sum and a count is. Machine A emits, for cluster 1, the sum 4 and the count 2. Machine B emits, for cluster 2, the sum 10 and the count 1. Add the sums, add the counts, divide once at the top: c1 is 4/2, which is 2, and c2 is 10/1, which is 10. Those are the two numbers a single machine produced a moment ago.

That figure is the previous one cut down the middle, with a reduce bolted underneath. Nothing else about the algorithm changed.

Now the names, which I have been holding back. The assignment step is embarrassingly parallel: the work splits with no communication across the splits. Pushing the centroids out to both machines is a broadcast. Folding the per-machine sums and counts into one pair per cluster is a tree-reduce, which meets partial results in pairs the way a knockout bracket does, so its depth grows with the logarithm of the worker count rather than with the count itself. A reduce whose result immediately broadcasts back out is an all-reduce.

Four terms, and they describe the two paragraphs above without adding anything to them.

Now make it a billion points

Another rung up, and the procedure does not change. There is just three orders of magnitude more of it.

Split the billion points across workers and pin them there. Each worker holds a partition, assigns its own partition against the current centroids, and emits per-cluster sums and counts. The partition itself does not travel, in any iteration.

Here is the arithmetic that makes the arrangement worth its complexity. Going out: a billion points times 8 dimensions times 8 bytes per float64 is 64 GB. Crossing the wire per iteration, with k set to 5: five centroids times 8 dimensions times 8 bytes is 320 bytes. The ratio between them is $2 \times 10^{8}$.

Coming back up, each worker emits k sums and k counts. Five sums of eight float64 is 320 bytes, five int64 counts is 40 bytes, so 360 bytes per worker per iteration. At a hundred workers that is roughly 36 KB climbing the reduce tree. Set against 64 GB of pinned points, both directions round to zero. A broadcast is cheap the way a memo is cheap: one small thing copied onto many desks, when the alternative on the table was moving the filing cabinets.

The reason the ratio is that lopsided is the shape of the per-iteration cost:

k and d are small and you chose them. n is the one you were handed, and it is the only term that grew.

One storage decision belongs here, because it is the same multiplication. Store the features as float32 rather than float64 and the constructed dataset goes from 64 GB to 32 GB. On a laptop that is a nicety. Across a cluster it decides whether a partition sits in a worker’s memory or spills to that worker’s disk, and spilling is the thing the cluster was supposed to avoid.

If the mechanism is what you came for, that is all of it, and this is a fine place to stop. The three sections below are the parts that bite once you run it.

The step that did not get cheaper

The loop distributes cleanly. Initialization does not, and that asymmetry is the main practical surprise in distributed k-means.

Classic k-means++ picks seeds one at a time. Each new centroid is drawn with probability proportional to its squared distance from the nearest seed already chosen, so the probabilities depend on the current seed set, so the data has to be scanned before the next one can be drawn. That is k sequential passes. On one machine it is k passes over an array. On a cluster each pass is a full distributed scan, and the whole stack of them happens before the first assignment step of the actual loop.

Line the two costs up. Seeding costs k distributed passes. The loop that follows costs i iterations, and i is usually a few tens. At k of 5000 and i of 30, seeding is about 167 times as many distributed passes as the clustering it was preparing for. Be careful what that ratio counts. It is passes, not floating-point work: a seeding pass scores every point against one new center, where a Lloyd iteration scores every point against all k, so the arithmetic per seeding pass is much lighter. On a cluster the pass is still the unit that hurts, because each one is a scheduled job over the whole dataset with its own stragglers to wait on.

That leaves a candidate pool bigger than k, and collapsing it is not free. Every point in the dataset has to be scored against the pool so that each candidate can be weighted by how many points it is nearest to. That is one more distributed pass, not zero. Only then is the weighted pool, which is small, reclustered down to k seeds on a single machine. So the honest count is a logarithmic number of oversampling passes plus one weighting pass, set against k sequential ones. Still the trade you want at large k, and still not the free lunch the summaries usually imply.

In practice the round count is smaller than the theory has to allow for. Spark MLlib defaults to two oversampling steps, which is a pragmatic setting rather than a reading of the bound.

from pyspark.ml.clustering import KMeans as SparkKMeans

# initMode picks the seeder. initSteps is the number of
# oversampling passes k-means|| makes before the weighting pass.
kmeans = SparkKMeans(k=5, initMode="k-means||", initSteps=2)

Two libraries that already do this

PySpark MLlib and Dask-ML both implement the data movement from the three sections above. Picking between them is an ecosystem question, and I would treat any per-iteration performance argument between the two with suspicion unless it arrives with the hardware attached.

Reach for Spark when the data already lives in the JVM ecosystem. Load it, assemble the feature column, fit. The inferSchema option is load-bearing: without it every column reads back as a string and VectorAssembler refuses the frame.

from pyspark.ml.clustering import KMeans as SparkKMeans
from pyspark.ml.feature import VectorAssembler

df = (spark.read
      .option("header", "true")
      .option("inferSchema", "true")
      .csv(data_path))
features = VectorAssembler(inputCols=numeric_cols, outputCol="features")
df_vec = features.transform(df)

model = SparkKMeans(k=5, maxIter=50, featuresCol="features").fit(df_vec)
predictions = model.transform(df_vec)

Reach for Dask when the workflow is already NumPy and pandas native. The API mirrors scikit-learn closely enough that the diff against single-machine code is a couple of lines. Both libraries made the same seeding choice, incidentally: dask_ml.cluster.KMeans defaults to init="k-means||" exactly as Spark does, so the previous section describes what either one is doing before your first iteration. I wrote up the delayed-graph side of Dask in an earlier post; this is the array side of the same library.

import dask.dataframe as dd
from dask_ml.cluster import KMeans as DaskKMeans

df = dd.read_csv(data_path)
X = df[numeric_cols].to_dask_array(lengths=True)

# init_max_iter is Dask-ML's equivalent of initSteps: the number of
# k-means|| oversampling rounds before the weighting pass.
kmeans = DaskKMeans(n_clusters=5, max_iter=100, init_max_iter=3)
kmeans.fit(X)
labels = kmeans.predict(X)

If the pipeline is JVM and the data is already in Spark, stay in Spark. If the pipeline is Python and the data is already in arrays, stay in Dask. The cost of crossing that boundary, which means serializing a large dataset out of one runtime and into the other, tends to exceed whatever difference was being chased across it.

When there is no dataset to fit

Everything so far assumes the data eventually lands somewhere and you fit over the whole of it. Drop that assumption and you have a stream, where batches arrive and each one is seen once.

The tool for that is partial_fit, which updates the centroids from one mini-batch and then forgets the batch. It treats data the way a running average treats a sensor feed. Keep the estimate, drop the reading.

There is a sharp edge here, because the obvious call does not exist:

>>> from dask_ml.cluster import KMeans   # dask-ml 1.4.0
>>> KMeans(n_clusters=3).partial_fit(X)
AttributeError: 'KMeans' object has no attribute 'partial_fit'

from sklearn.cluster import MiniBatchKMeans

def streaming_kmeans(data_stream, k=3):
    model = MiniBatchKMeans(n_clusters=k, random_state=42)
    for batch in data_stream:
        model.partial_fit(batch)   # this batch moves the centroids, then is dropped
    return model

When you own that loop, wrapping it in dask_ml.wrappers.Incremental adds nothing. Incremental is for the other case, where the data is already one large Dask array and you want partial_fit mapped across its blocks rather than written out by hand.

from sklearn.cluster import MiniBatchKMeans
from dask_ml.wrappers import Incremental   # dask-ml 1.4.0

est = Incremental(MiniBatchKMeans(n_clusters=3, random_state=42))
est.fit(X)   # X is a Dask array; Incremental calls partial_fit block by block

This is the regime where “it fits across the cluster” stops meaning anything, because there is no fixed dataset, no fixed size, only an arrival rate.

What the ladder was for

Read backwards, the post is one move shown at three sizes: one machine, two machines, a thousand. Find the step whose work splits per record. Run it where the records already sit. Ship a summary rather than the records.

Data-parallel SGD has the same shape with a different payload. Each worker computes gradients on its own shard, an all-reduce sums them, the summed gradient goes back out, and the next step happens. What differs is the size of the thing crossing the wire. For k-means it is k times d floats, and it stays that size whether the cluster holds a million points or a billion. For SGD it is the model’s full gradient vector, which grows with the model, which is why a framework like Horovod moves gradients with a ring all-reduce, where each worker only ever talks to its two neighbours and the bytes each one sends stay flat as the cluster grows. Nobody has ever needed that for five centroids. K-means is the version of the pattern where the model is small enough to watch.

Two things I left out on purpose. There is nothing here about choosing k: the elbow sweep, the silhouette score, and the automated versions of both are their own subject, and they mostly do not change under distribution, beyond the obvious point that a sweep multiplies everything above by the number of values tried. There is also nothing about skewed partitions, where one worker ends up holding most of a cluster’s points and the reduce sits waiting on it. That one I would want measurements for before writing about it, and I do not have any.

What is left to pay for is the coordination. Workers have to agree on who holds what, the broadcast has to land, and the reduce has to complete before the next iteration can start. That is a fixed cost, paid before any clustering happens. Below some size it is a cost with nothing to show for it, and a single-node MiniBatchKMeans will finish first, though that one is trading inertia for speed and is not solving quite the same problem. The gap between a constructed 64 GB and a 16 GB laptop is what pays for the machinery. Underneath that gap, the machinery is the problem.

Part one: parallelizing the feature loop Scalable K-Means++, Bahmani et al. Spark MLlib clustering