k-means: the building block

Both IVF and quantization rest on one algorithm: k-means clustering. It finds k representative points (centroids) that summarize a cloud of vectors. Get this and the rest of the book is just two clever ways of using centroids.

The goal

Given many vectors, find k centroids so that every vector is close to its nearest centroid. Those centroids become:

the region centers in IVF (each vector belongs to its nearest region), and
the palette in quantization (each vector is replaced by its nearest centroid's id).

Lloyd's algorithm (the standard k-means)

It alternates two steps until things stop moving:

Assign — put each vector in the group of its nearest centroid.
Update — move each centroid to the average of its group.

Repeat. Each round can only reduce the total distance, so it converges.

start: random centroids        assign points        move centroids to means
   x   x      ●                 x   x  |  ●           x  x        ●
 x   x          ●      ─►     x   x    |    ●   ─►   x  ●            ●
   x    x   x                   x  x x | (nearest)     x  x  x
                                                    ...repeat until stable

The code

import numpy as np


def _assign(X, C):
    """Index of the nearest centroid for each row of X (squared L2).
    Uses ||x-c||^2 = ||x||^2 - 2 x.c + ||c||^2 to compute all distances at once."""
    d = (X ** 2).sum(1)[:, None] - 2.0 * X @ C.T + (C ** 2).sum(1)[None, :]
    return d.argmin(1)


def kmeans(X, k, iters=25, seed=0):
    """Lloyd's k-means. Returns (centroids, assignments)."""
    X = np.asarray(X, dtype=np.float64)
    n = X.shape[0]
    k = min(k, n)
    rng = np.random.default_rng(seed)
    C = X[rng.choice(n, size=k, replace=False)].copy()   # init: random points
    for _ in range(iters):
        a = _assign(X, C)                                # 1. assign
        newC = C.copy()
        for j in range(k):                              # 2. update
            members = X[a == j]
            if len(members):
                newC[j] = members.mean(0)
        if np.allclose(newC, C):                        # converged?
            C = newC
            break
        C = newC
    return C, _assign(X, C)

Reading it

_assign computes, for every vector, the distance to every centroid in one vectorized expression, then picks the closest. This is the workhorse — IVF and PQ both call it.
kmeans starts from k random data points, then loops assign→update until the centroids stop changing (or iters is hit).

See it work

Six points forming two obvious clusters — one near the origin, one near $(5,5)$:

import numpy as np
from ivfpq import kmeans

X = np.array([[0., 0.], [0.2, 0.1], [0.1, 0.2],     # cluster A
              [5., 5.], [5.1, 4.9], [4.9, 5.1]])     # cluster B
C, a = kmeans(X, k=2, seed=0)
print("centroids:\n", np.round(C, 2))
print("assignments:", a.tolist())

Output:

centroids:
 [[5.  5. ]
 [0.1 0.1]]
assignments: [1, 1, 1, 0, 0, 0]

k-means found exactly the two cluster centers — $(5,5)$ and $(0.1,0.1)$ — and assigned the first three points to centroid 1 and the last three to centroid 0. 🎯 Those two centroids now summarize all six points.

How centroids power everything next

IVF: treat each centroid as a bucket. Every vector is filed under its nearest centroid; at query time we only open the buckets near the query.
Quantization: treat the centroids as a codebook. Store, per vector, just the id of its nearest centroid (one small integer) instead of the full vector.

Same algorithm, two superpowers. Let's build IVF first. 👉

IVF & Product Quantization from Scratch