K‐means clustering

k-means clustering is an important operation in Faiss because it is at the foundation of quantization operations and IVF datasets.

Why k-means?

Faiss relies on the k-means clustering algorithm to do quantization in most cases.

k-means assigns vectors $\lbrace x_1,\dots,x_n\rbrace \in \mathbb{R}^d$ to a set of $k$ centroids $\lbrace c_1,\dots,c_k \rbrace$, by minimizing the quantization error in the mean squared error sense:

$$\mathrm{MSE} = \frac{1}{n}\sum_{i=n} \lVert x_i - c_{a_i} \rVert ^2$$

where the assignment of the vectors to centroids is

$$a_i = \mathrm{argmin}_{j=1..k} \lVert x_i - c_j \rVert ^2$$

The k-means algorithm starts from an initial (random) assignment and iterates two update steps:

1. update each centroid as the center of mass of the vectors assigned to it:

$$c_j := \frac{1}{|A_j|} \sum_{x\in A_j} x ~~~~ \textrm{where} ~~ A_j = \lbrace i=1..k \mathrm{st.} a_i = j\rbrace ~~~~ \forall j=1..k$$

2. re-compute assignment using the argmin above.

There are additional heuristics to avoid that some centroids collapse (no point is assigned to them) This algorithm is guaranteed to converge because both steps decrease the MSE and the number of possible assignments is finite. However, it is not guaranteed to reach a global optimum. Computing the global minimizer of the MSE has a complexity that is exponential in $n$, and thus NP-hard.

Note that the quantization error is defined in terms of squared L2 distance because this enables the center of mass update of the centroid (other distances not minimized as easily, but the convexity of the distance function helps).

k-means to make compact clusters

k-means attempts to minimize the distance to the cluster centroids. However, a side-effect of this is that is minimizes the squared pairwise distances within clusters. Indeed, for a cluster with $n$ vectors $\lbrace x_1,\dots,x_n\rbrace \in \mathbb{R}^d$, the pairwise distances are given by

$$ T = \sum_{i=1}^n \sum_{j=1}^n \lVert x_i - x_j \rVert^2 $$

Given the center of mass of the cluster ($\bar{x} = \frac{1}{n}\sum_{i=1}^n x_i$), T can be decomposed into

$$ T = \sum_{i=1}^n \sum_{j=1}^n \lVert (x_i - \bar{x}) - (x_j - \bar{x}) \rVert^2 = \sum_{i=1}^n \sum_{j=1}^n \left( \lVert x_i - \bar{x} \rVert^2 + \lVert x_j - \bar{x} \rVert^2 - 2 \langle x_i - \bar{x}, x_j - \bar{x} \rangle \right) $$

Reorganizing a bit, this gives:

$$ T = 2n \sum_{i=1}^n \lVert x_i - \bar{x} \rVert^2 - 2 \sum_{i=1}^n \langle x_i - \bar{x}, \left( \sum_{j=1}^n (x_j - \bar{x}) \right) \rangle $$

The second term vanishes, since $\sum_{j=1}^n (x_j - \bar{x}) = \sum_{j=1}^n (x_j - \frac{1}{n}\sum_{i=1}^n x_i) = 0$, so

$$ T = 2n^2 \times \mathrm{MSE} $$

Therefore, minimizing the sum of squared distances within a cluster is equivalent to minimizing the mean squared distance to the centroid $\bar{x}$, which is the k-means objective.

k-means for the cosine similarity

For simplicity of notation, we consider that vectors and centroids are L2 normalized, then the cosine similarity is a dot product. To best represent vectors, we want to maximize:

$$\mathrm{CS} = \frac{1}{n}\sum_{i=n} \langle x_i, c_{a_i} \rangle$$

Fixing the assignment and updating centroid $c_j$ (k-means step 1) is done via

$$\max_{c\in\mathbb{R}^d~~\mathrm{st.}\lVert c\rVert = 1} \sum_{i\in A_j} \langle x_i, c \rangle = \max_{c\in\mathbb{R}^d\mathrm{st.}~~\lVert c\rVert = 1} \langle \sum_{i\in A_j} x_i, c \rangle, ~~~~\textrm{Which is solved by}~~~~c_j=\frac{\sum_{i\in A_j} x_i}{\lVert \sum_{i\in A_j} x_i \rVert}$$

Recomputing the assignment (k-means step 2) is simply

$$a_i = \mathrm{argmax}_{j=1..k} \langle x_i, c_j \rangle$$

These two steps are guaranteed to converge to a local optimum. This justifies the "normalized k-means" (or spherical k-means) in Faiss for the cosine similarity.

Simple Python interface

Faiss provides an efficient k-means implementation. Cluster a set of vectors stored in a given 2-D tensor x is done as follows:

ncentroids = 1024
niter = 20
verbose = True
d = x.shape[1]
kmeans = faiss.Kmeans(d, ncentroids, niter=niter, verbose=verbose)
kmeans.train(x)

The resulting centroids are in kmeans.centroids.

The values of the objective function (total square error in case of k-means) along iterations is stored in the variable kmeans.obj, and more extensive statistics in kmeans.iteration_stats.

To run it on GPU, add the option gpu=True to the Kmeans constructor. This will use all available GPUs on the machine.

Additional options

The Kmeans object is mainly a layer of the C++ Clustering object (see below), and all fields of that object can be set via the constructor. The fields include:

• nredo: run the clustering this number of times, and keep the best centroids (selected according to clustering objective)

• verbose: make clustering more verbose

• spherical: perform spherical k-means -- the centroids are L2 normalized after each iteration

• int_centroids: round centroids coordinates to integer

• update_index: re-train index after each iteration?

• min_points_per_centroid / max_points_per_centroid: below, you get a warning, above the training set is subsampled

• seed: seed for the random number generator

Assignment

To compute the mapping from a set of vectors x to the cluster centroids after kmeans has finished training, use:

D, I = kmeans.index.search(x, 1)

This will return the nearest centroid for each line vector in x in I. D contains the squared L2 distances.

For the reverse operation, eg. to find the 15 nearest points in x to the computed centroids, a new index must be used:

index = faiss.IndexFlatL2 (d)
index.add (x)
D, I = index.search (kmeans.centroids, 15)

I of size (ncentroids, 15) contains the nearest neighbors for each centroid.

Clustering on the GPU

Clustering on one or several GPUs can be done via the gpu=True (use all gpus) or gpu=3 (use 3 gpus) constructor option in the KMeans object.

Lower level C++ interface

The C++ interface is based on the Clustering object. It maintains the set of centroids, and offers a training method to perform the k-means iterations.

See this GPU k-means wiki on how to call the GPU k-means from C++.

Alternative indexes to perform assignment

The most expensive k-means operation is the assignment of the training vectors to the current set of centroids, which is a nearest-neighbor search problem. Therefore, the Clustering::train() method takes an index as argument. This index can be a CPU index (like the basic IndexFlat) or a GPU index or another approximate index.

In the experiment below (from bench_hnsw.py), the assignment step in k-means is performed with a HNSW index instead of a flat assignment:

  performing kmeans on the sift1M vectors (baseline)
  Clustering 1000000 points in 128D to 16384 clusters, redo 1 times, 10 iterations
    Preprocessing in 0.17 s
    Iteration 9 (612.53 s, search 612.18 s): objective=3.85228e+10 imbalance=1.235 nsplit=0

  performing kmeans on the sift1M using HNSW assignment
  Clustering 1000000 points in 128D to 16384 clusters, redo 1 times, 10 iterations
    Preprocessing in 0.17 s
    Iteration 9 (74.63 s, search 73.46 s): objective=3.85232e+10 imbalance=1.234 nsplit=0

ie. the clustering is 8.2x times faster than with an exhaustive assignment, without impacting the objective value or increasing the number of empty clusters (nsplit) or imbalance factor, both are signs of unhealthy clusters.

2-level clustering

This is one very simple approach was suggested by Harsha Simhadri, from Microsoft, who organizes the billion-scale ANN challenge. It simply consists in doing a two-level hierachical clustering to obtain the centroids. However, the centroids are then considered in a full table (ie. there is no hierarchical routing at assignment time).

This is a cheap and alternative to the full clustering when distributed clustering and/or GPUs are not available. It is implemented in the Faiss contrib as two_level_clustering. The accuracy impact is very limited, see eg. this comparison or coarse quantizers.

Pure Python implementation of k-means and distributed k-means

The pure Python implementation of k-means abstracts the dataset as a class DatasetAssign that offers the following operations:

• sample a set of points from the dataset (get_subset)

• assign the dataset points to a given set of centroids (perform_search)

• accumulate the per-cluster training point sums (assign_to)

The function kmeans uses this DatasetAssign to perform the iterations.

This allows, for example, to cluster sparse vectors with a DatasetAssignSparse.

The k-means clustering is easy to distribute over several machines (or GPUs). The training set is sharded the machines, and a central machine keeps track of the centroids via a DatasetAssign object.

This has been implemented here to cluster 500M points to 10M centroids.