K-Means with Batched Restarts

Run many k-means restarts at the same time, as one batched tensor program.

Lloyd's algorithm finds a local minimum of the inertia, and which minimum it finds depends on the start. The usual fix is to run it from several random starts and keep the best result. A loop over restarts works, but on a GPU you want one batch: put the restarts on a leading axis and let every operation handle all of them at once.

C:        (R, k, d)   centroids of every restart
X:        (n, d)      the data, shared by all restarts
dist:     (R, n, k)   torch.cdist of the expanded X against C
labels:   (R, n)      nearest centroid for each point, in each restart

Each iteration assigns every point to its nearest centroid, then moves each centroid to the mean of its points. Do the mean with scatter_add_ (or index_add_) along the cluster axis, so that no Python loop runs over points or clusters. At the end, compute the inertia of each restart and return the centroids and labels of the best one.

inertia = sum over points of the squared distance to the assigned centroid

Function signature

def batched_kmeans(X, k, n_restarts=10, n_iters=50, seed=0):

Example

Input

X = torch.tensor([[0., 0.], [0., 1.], [10., 10.], [10., 11.]])
centroids, labels = batched_kmeans(X, k=2)

Output

centroids  ->  tensor([[ 0.0,  0.5],
                       [10.0, 10.5]])     (row order can differ)
labels     ->  tensor([0, 0, 1, 1])       (ids can be swapped)

Reasoning

The two tight pairs are the two clusters. Each centroid is the mean of its pair, and the inertia is 4 * 0.5 ** 2 = 1.0, the best possible.

Constraints and notes

• Initialise each restart with k distinct rows of X, chosen at random with a torch.Generator seeded by seed. The same seed must give the same result.
• A cluster can lose all its points. Keep its old centroid in that case. A plain sums / counts divides by zero and returns nan.
• Stop early when no centroid moves. Return labels that match the returned centroids, that is, each label is the nearest centroid.
• Return centroids of shape (k, d) in the dtype of X, and labels of shape (n,) and dtype int64.
• Graded on the result: the inertia must be within 1% of scikit-learn's KMeans with 10 restarts.