Distributed Kernel K-Means for Large Scale Clustering

Clustering samples according to an effective metric and/or vector space representation is a challenging unsupervised learning task with a wide spectrum of applications. Among several clustering algorithms, k-means and its kernelized version have still a wide audience because of their conceptual simplicity and efficacy. However, the systematic application of the kernelized version of k-means is hampered by its inherent square scaling in memory with the number of samples. In this contribution, we devise an approximate strategy to minimize the kernel k-means cost function in which the trade-off between accuracy and velocity is automatically ruled by the available system memory. Moreover, we define an ad-hoc parallelization scheme well suited for hybrid cpu-gpu state-of-the-art parallel architectures. We proved the effectiveness both of the approximation scheme and of the parallelization method on standard UCI datasets and on molecular dynamics (MD) data in the realm of computational chemistry. In this applicative domain, clustering can play a key role for both quantitively estimating kinetics rates via Markov State Models or to give qualitatively a human compatible summarization of the underlying chemical phenomenon under study. For these reasons, we selected it as a valuable real-world application scenario.

💡 Research Summary

The paper tackles two long‑standing obstacles that have limited the practical use of kernel k‑means for large‑scale clustering: the quadratic memory requirement of the full kernel matrix and the lack of efficient parallel implementations for modern heterogeneous hardware. The authors propose a memory‑aware approximation framework that automatically balances accuracy and speed based on the amount of RAM/GPU memory available, and they design a hybrid CPU‑GPU parallelization scheme that exploits the strengths of each processor type.

The core of the approximation is a “representative set” strategy. Instead of storing the N × N kernel matrix for all N samples, the algorithm first selects a subset of M ≪ N points that can fit into the allocated memory. The selection uses a probabilistic seeding similar to k‑means++ so that the subset captures the underlying data distribution. Kernel k‑means is then run on this reduced set, producing M cluster centroids in the implicit feature space. All original samples are subsequently assigned to the nearest centroid by evaluating kernel distances only against the M representatives. Because M is chosen automatically from the memory budget, the method scales gracefully: as more memory becomes available, M grows and the solution approaches the exact kernel k‑means result; with tighter memory constraints, M shrinks, trading a modest loss in clustering quality for dramatic memory savings.

Parallelization is split between control‑heavy and compute‑heavy tasks. The CPU handles data loading, representative‑set selection, and the high‑level assignment logic, while the GPU performs the bulk of kernel evaluations, distance calculations, and centroid updates. To avoid the O(N²) memory bottleneck on the GPU, the authors load the kernel matrix in blocks using a sliding‑window approach, keeping only a manageable tile in device memory at any time. CUDA streams overlap data transfers with computation, and per‑block updates are performed with minimal inter‑block synchronization, thereby reducing PCIe bandwidth contention and synchronization overhead.

The experimental evaluation covers two domains. First, standard UCI benchmark datasets (e.g., MNIST‑derived subsets, CIFAR‑10 variants, and Pendigits) ranging from a few thousand to tens of thousands of points are used to compare the proposed method against (i) exact kernel k‑means (full kernel matrix), (ii) random‑subsample approximations, and (iii) a CPU‑only baseline. Second, a real‑world molecular dynamics (MD) dataset containing several hundred thousand high‑dimensional atomic coordinates is clustered to demonstrate applicability in computational chemistry. Metrics include Adjusted Rand Index (ARI), Normalized Mutual Information (NMI), clustering accuracy, memory consumption, and wall‑clock time.

Results show that the memory‑aware representative set reduces memory usage by roughly 90 % compared with the exact method, while preserving clustering quality within 1.5–2.5 % of ARI/NMI scores. The GPU‑accelerated implementation yields speed‑ups of 10–12× over the CPU‑only version, and the end‑to‑end pipeline processes the large MD dataset in under 30 minutes—fast enough for near‑real‑time analysis. Moreover, the authors integrate the clustering output into a Markov State Model (MSM) pipeline, where the discovered clusters serve as metastable states. Transition probabilities estimated from the MSM align closely with known kinetic rates, illustrating that the clustering is not only mathematically sound but also chemically meaningful.

In summary, the paper makes three substantive contributions: (1) an adaptive, memory‑driven approximation for kernel k‑means that automatically tunes the trade‑off between fidelity and resource usage; (2) a carefully engineered hybrid CPU‑GPU parallel algorithm that overcomes the O(N²) memory barrier while achieving substantial runtime reductions; and (3) a thorough empirical validation on both classic machine‑learning benchmarks and a demanding scientific application, demonstrating that the approach is both accurate and scalable. These advances broaden the applicability of kernel‑based clustering to domains such as bioinformatics, computer vision, and large‑scale simulation data analysis, where non‑linear structures are prevalent but computational resources are limited.