A Super Fast K-means for Indexing Vector Embeddings

Abstract: We present SuperKMeans: a k-means variant designed for clustering collections of high-dimensional vector embeddings. SuperKMeans' clustering is up to 7x faster than FAISS and Scikit-Learn on modern CPUs and up to 4x faster than cuVS on GPUs (Figure 1), while maintaining the quality of the resulting centroids for vector similarity search tasks. SuperKMeans acceleration comes from reducing data-access and compute overhead by reliably and efficiently pruning dimensions that are not needed to assign a vector to a centroid. Furthermore, we present Early Termination by Recall, a novel mechanism that early-terminates k-means when the quality of the centroids for retrieval tasks stops improving across iterations. In practice, this further reduces runtimes without compromising retrieval quality. We open-source our implementation at https://github.com/cwida/SuperKMeans

• The paper introduces SuperKMeans, a fast k-means algorithm that integrates dimension pruning and GEMM operations for high-dimensional vector embeddings.
• It achieves up to 7x speedups on CPU and 4x on GPU by pruning over 97% of centroids and employing recall-driven early termination.
• Empirical evaluations show scalability over diverse architectures and datasets, ensuring minimal quality degradation for vector search indexing.

SuperKMeans: Accelerating K-Means for High-Dimensional Vector Embedding Indexing

Introduction and Motivation

The scalability and efficiency of clustering high-dimensional vector embeddings are essential for contemporary vector search systems, particularly those underpinning large-scale information retrieval, recommendation engines, and retrieval-augmented generation (RAG). Traditional $k$-means implementations, including those found in state-of-the-art libraries such as FAISS and Scikit-Learn, remain limited by their reliance on a vanilla Lloyd's algorithm, with primary optimizations tethered to highly tuned BLAS GEMM routines. However, the need for large $k$ and high $d$ values in partition-based indexes (e.g., IVF) exposes $k$-means to severe compute bottlenecks.

SuperKMeans addresses these bottlenecks by tightly integrating dimension pruning via adaptive sampling and progressive thresholding, interleaved with GEMM operations and specialized memory layouts. This results in substantial speedups—up to 7x over CPU baselines and 4x over GPU baselines—while maintaining clustering quality and index suitability for vector similarity search (VSS).

Figure 1: Clustering performance for all competitors in OpenAI's 1M vectors of 1536 dimensions; SuperKMeans yields up to 7x (CPU) and 4x (GPU) speedups as $k$ grows.

Algorithmic Contributions

SuperKMeans' acceleration centers on two-phase iteration within the assignment step of Lloyd's algorithm:

• Phase 1 (GEMM): Compute partial distances using the leading $d'$ dimensions via GEMM, optimizing cache locality and SIMD utilization.
• Phase 2 (PRUNING): Employ Adaptive Sampling (ADSampling) and memory-resident centroid layouts (PDX) to prune over 97% of centroids before full-dimensional computations. Progressive pruning with block decomposition further reduces unnecessary arithmetic.

  Figure 2: SuperKMeans iteration phases: first, GEMM for partial distances; second, PRUNING for full assignment via dimension-wise selective evaluation.

Early Termination by Recall (ETR) is introduced, stopping iterations once further improvements in VSS recall plateau, eliminating guesswork for iteration count and reducing redundant computation. ETR can be transferred to any iterative clustering pipeline.

Empirical Evaluation

Clustering Time and Quality

SuperKMeans' superiority is demonstrated across ten modern embedding datasets and multiple CPU microarchitectures—Intel Granite Rapids, AMD Zen variants, AWS Graviton, and Apple M4—using dataset-specific BLAS kernels. With $k = 4\sqrt{N}$, SuperKMeans achieves 3-4x speedups in high-dimensional settings and matches FAISS and Scikit-Learn in recall and cluster balance.

Figure 3: Clustering time (CPU) for various competitors across multiple datasets; SuperKMeans excels in $d > 384$ settings.

Scalability with k

As $k$ increases to 100,000, SuperKMeans maintains or increases its speedup factor relative to competitors, confirming its suitability for large-scale vector indexes. At low $k$, performance converges due to diminished benefits of pruning (all centroids fit within cache).

Figure 4: Speedup vs k on CPU; SuperKMeans accelerates increasingly with larger k.

Impact of Sampling and Iterations

SuperKMeans maintains high recall even when clustering with 20–30% of the data, and cluster balance stabilizes at these sampling rates. ETR generally terminates after 5–10 iterations, minimizing computational cost with negligible recall degradation.

Figure 5: ETR triggers fewer iterations with minor impact on VSS recall; cluster balance stabilizes rapidly.

Figure 6: Sampling effects: minimal recall loss, moderate increase in vectors explored at low sample fractions.

Multithreading and Microarchitecture Portability

SuperKMeans scales nearly linearly with thread count and demonstrates consistent speedups across all SIMD instruction sets and CPU architectures, outperforming both FAISS and Scikit-Learn, especially when register sizes are limited (Graviton 4, Apple M4).

Figure 7: SuperKMeans speedups on diverse CPU microarchitectures, demonstrating portability.

Initialization and Ablation Results

Random centroid initialization matches or surpasses $k$-means++ in speed and index quality for VSS. Progressive pruning, PDXification, and threshold adjustment combine for maximal speedup; omitting these optimizations incurs substantial slowdowns.

Figure 8: $k$-means++ initialization is orders of magnitude slower and not beneficial for index quality.

GPU Implementation

On NVIDIA H100, SuperKMeans retains its two-phase architecture, utilizing cuBLAS for GEMM and custom CUDA kernels for block-wise pruning. Aggressive parallelization and register-level memory overlapping mitigate occupancy bottlenecks. Speedups against FAISS and cuVS are most pronounced at large k; preprocessing cost on CPU constrains speedup for small k.

Figure 9: SuperKMeans achieves up to 4x GPU speedup as k increases; preprocessing is amortized for large k.

Hierarchical SuperKMeans

Hierarchical $k$-means, where SuperKMeans is deployed in both meso- and fine-clustering phases, achieves even greater speedups—up to 250x over FAISS—enabling blazing-fast clustering for very large collections (>10M vectors). The trade-off is minor recall degradation (<3%) but improved cluster balance and reduced vectors explored in VSS.

Figure 10: Hierarchical SuperKMeans attains unprecedented 250x clustering speedup vs FAISS.

Theoretical and Practical Implications

SuperKMeans is applicable to any variant requiring fast assignment over large $k$ and $d$. ETR simplifies hyperparameter tuning and can benefit iterative quantization and refinement pipelines, e.g., Product Quantization and hybrid indexes. The principal insight is that most centroids are functionally irrelevant for assignment in high-dimensional spaces, enabling pragmatic acceleration by design.

Furthermore, the core pruning motif and memory layout advances (PDX) also transfer to VSS, drastically improving query speed and footprint relative to traditional IVF and graph-based (HNSW) approaches.

Conclusion

SuperKMeans decisively advances the practical speed and scalability of $k$-means clustering for high-dimensional vector embedding collections used in modern vector search engines. By combining interleaved GEMM with dimension-adaptive pruning, and by plugging in recall-driven early termination, SuperKMeans accomplishes superior runtime without quality loss, on both CPUs and GPUs. Its modular design enables integration with hierarchical and minibatch variants, as well as direct application to advanced quantization and indexing methods. The findings indicate that state-of-the-art vector search systems should migrate toward SuperKMeans-based index pipelines to unlock efficient, scalable embedding clustering.

Future work includes integration with dimensionality-reduced embeddings, further GPU kernel optimization (e.g., Flash-KMeans), and expanding SuperKMeans to additional iterative index construction paradigms.

Figure 11: PDX pruning, originally motivated by SuperKMeans, makes IVF competitive with HNSW, reducing index build time and storage footprint.