Faiss 1.16.0: Distributed e-means training, sample-based clustering

1. Train together using summaries without gathering the data

e-means processes data in small batches and updates centroids, the representative positions of clusters. In the new distributed implementation, each process, an independently running task, holds all centroids. Processes combine summaries such as vector sums and counts to update centroids and aggregate the overall training error.

Users create a PyTorch distributed communication group and assign non-overlapping rows to each process. Instead of passing the entire array at once, they orchestrate training initialization, data accumulation, centroid updates, pass completion, and finalization. CPUs use Gloo, and the original demo shows how to run two processes on GPUs and CPUs.

2. Reduce data reads and training with a sample cap

The new max_points_per_centroid option sets the maximum number of training rows per centroid. If the total row count exceeds the cap multiplied by the centroid count, a random sample is drawn using a seed, a value that makes random selection reproducible. The default is None, which trains on all rows as before.

The batch size, the number of data items processed together, and training-step calculations also adjust to the sample size. For NumPy input, only selected rows are read, so memory-mapped datasets can be trained without loading the entire dataset into memory. Memory mapping accesses a file through memory addresses. The following example is reproduced unchanged from the proposed sample-limit change.

km = Emeans(d, k, max_points_per_centroid=256)
centroids, labels = km.train(x)   # labels[i] == -1 for rows outside the sample

The returned label array retains the input row count, with -1 instead of a cluster number for rows outside the sample. Thus, even with a smaller training sample, the output array preserves its correspondence to the original input rows.

3. Benchmarks comparing time to reach the same error

This release also includes scripts that remeasure results from the e-means paper. The company's results below show the epochs and speedups for e-means to reach the test error of Flash-KMeans after 25 iterations on one NVIDIA H100 80GB. An epoch is one pass through the training data.

Test error is the mean squared distance between held-out vectors and their nearest centroids. For each dataset and seed, both methods used the same GPU and initial centroids. Training time was measured after warm-up, excluding data loading and error evaluation.

Epochs are summarized by sorting the times to target across seeds and taking the larger of the two middle values. Speedup is the median of per-seed time ratios, the middle value when sorted. The SPACEV and SimSearchNet++ reference implementations were adjusted to compute squared vector norms in float32.

Table: e-means training effort and speedup at the Flash-KMeans reference error

Conditions: H100 80GB, PyTorch 2.14, Triton 3.8, CUTLASS DSL 4.7.1. Flash-KMeans 0.3.1, float16, 25 iterations, seeds 10~19. SPACEV and SimSearchNet++ reference implementations use float32 squared norms, which express squared vector magnitudes. e-means uses the default recipe, CuTe kernels, bfloat16 assignment, and per-batch conversion of float32 input. Includes a CuTe cache fix for information temporarily stored for reuse.

Units: epochs, x

1M, 10M, and 100M denote 1 million, 10 million, and 100 million rows, respectively. The time ratio divides the reference method's training time by the e-means time for each dataset, so larger values mean reaching the same error faster.

According to the company's results, a separate input path with SimSearchNet++ 100M preconverted to bfloat16 achieved 2.6 epochs and a 19.9x speedup. Both the table and this separate result measure clustering training on a single GPU.

4. Less data processing, more operational responsibility

A sample cap reduces the data read and trained on, while row distribution splits training data across devices. In distributed setups, users manage processes and data allocation and bear the communication costs between centroid replicas. Since each device retains the centroids, configuration choices must account for centroid memory separately from data distribution.

The sample cap applies to the array-based implementation; specifying it in distributed or centroid-sharded implementations raises an error. Row-distributed training requires the global batch, including the last batch, and the row count per epoch to be divisible by the process count. A global batch is processed jointly by all processes. The table's timings include a CuTe kernel cache fix being integrated separately. Without that fix, partial epochs ending before a full pass may take longer and yield lower speedups.

Faiss 1.16.0 lets users choose between distributing training rows across devices and reducing them through sampling, with the performance table serving as a single-GPU baseline.