011 — consolidated benchmark: initialization, Lloyd, and full fit

This page gathers every current result for the five-task mix in one place:

1. initialization only, before any Lloyd iteration;
2. Lloyd only, 12 iterations from the same random data rows for everyone;
3. full fit, initialization followed by 12 Lloyd iterations.

results.json holds 570 rows. 195 are full-fit and Lloyd-only rows from 008, excluding its retired stochastic-relocation variants. 180 are initializer rows from 009 and 010. 195 are new rows: Lloyd-only and full fits with the array-API farthest-point relocation, and full fits started from scalable k-means++ for the array API and cuML. ablation.json holds the 60-row scalable-init ablation explained in the last section.

Setup

• Tasks: k3-d1000 (950000 Gaussian-noise samples, 1000 features), k17-d180 (2100000 by 180), k100-d32 (1500000 by 32), k550-d6 (525000 by 6), and k3000-d2 (135000 quantized samples from 300 blobs in 2 dimensions). The task name gives the number of clusters and features.
• Contenders: the eager Python array-API implementation on PyTorch CPU (labelled Array API - Torch CPU), PyTorch CUDA (Array API - Torch CUDA), and CuPy (Array API - CuPy); stock scikit-learn 1.7.2; scikit-learn 1.10.dev0 with PRs 23 and 26; and cuML 26.10 on an NVIDIA L4.
• Initializers: greedy k-means++, scalable k-means++ (k-means||), and random data rows. Array-API rows use scikit-learn's farthest-point empty-center relocation unless labelled "keep empty".
• Float32, TF32 disabled, max_iter=12, tol=1e-8, seeds 0 to 2, 4 CPU threads, and a hard 60-second limit per worker process. All values are seed means; figures add one sample standard deviation.
• Figures use one color per implementation; hatching marks scalable initialization and dots mark random rows. Random initializers appear only on quality figures because they are not timing contenders. Quality is mean inertia divided by the best mean for the same task and section, with a dotted line at 1. Duration and memory use a logarithmic axis because their values span more than a factor of ten; every quality ratio stays below a factor of three, so those axes stay linear.

Initialization only

Initializer quality is the full-data mean inertia of the returned centers divided by the best initializer mean on each task. That best is cuML scalable on every task, for reasons analyzed below.

PyTorch CUDA greedy k-means++ is within a factor of 1.7 of cuML's greedy initializer. The array-API scalable initializer costs 0.9 to 2.6 s on PyTorch CUDA, more than its greedy counterpart on every task. cuML scalable is faster than cuML greedy from 17 to 550 clusters and ties it at 3000. Torch CPU scalable initialization needs 36 and 47 s at 550 and 3000 clusters.

The array-API row-norm feature tile dominates initializer memory on the two widest tasks. cuML and scikit-learn measure memory differently (nvidia-smi and VmHWM) and stay at or below 172 MiB.

All row-selecting initializers land within a few percent of each other, and array-API scalable improves on greedy by 2 to 5% at 17 and 100 clusters. cuML scalable is about 43% lower on the 3-cluster task and 35% lower on the 17-cluster task because it refines its candidate set with weighted Lloyd before returning.

Lloyd only from shared random rows

Every implementation starts from the same k random rows per seed, so this section isolates the Lloyd kernel and its empty-center policy. Its quality table divides each mean inertia by the best Lloyd-only mean for that task. Full-fit results do not contribute to this denominator.

PyTorch CUDA is the fastest Lloyd implementation at 3 and 17 clusters, for example 0.77 s against 6.85 s for cuML on the 1000-feature task. cuML is fastest from 100 clusters up, at 0.54, 0.35, and 0.32 s. The array-API tiles keep the peak increment near 0.9 GiB on the widest task, where scikit-learn peaks at 7.1 GiB and cuML at 3.8 GiB.

Implementations agree to the shown precision on the four tasks without empty clusters. Empty clusters only occur on the 3000-cluster task, with 46 to 55 empty-center events per run. There, farthest-point relocation sets the Lloyd-only reference on Array API - Torch CPU. The other array-API backends are within 0.002%, and scikit-learn is within 0.072%. Keeping empty centers is about 6.96% worse; cuML ends 6.99% above the reference, consistent with keeping its empty centers in place.

Full fit

PyTorch CUDA with greedy k-means++ is the fastest full fit at 3 and 17 clusters, at 0.94 and 1.25 s against 7.18 and 2.52 s for cuML. cuML is fastest from 100 clusters up, and cuML scalable improves its own greedy fit to 0.73 s at 100 clusters and 0.63 s at 550 clusters. The array-API scalable initializer adds 0.5 to 1.3 s to PyTorch CUDA fits and does not improve their final objective. Torch CPU exceeds the 60-second limit at 3000 clusters with either initializer and at 550 clusters with scalable initialization.

Twelve Lloyd iterations erase most initialization differences: every configuration but one ends within 0.4% of the best. Neither k-means++ variant produced an empty cluster, so keep-empty and farthest rows run the same computation. Their small differences on PyTorch CUDA come from rows measured in separate processes; seeded PyTorch CUDA draws are not bitwise reproducible across processes, while the CuPy rows match exactly.

The exception is cuML scalable at 3000 clusters: it starts from the best initializer objective but ends 5.6% above the best fit. Its centers are neither empty nor duplicated. With 300 iterations cuML converges to 0.01903 to 0.01916 from scalable starts against 0.01807 to 0.01811 from greedy starts, so this is a worse local optimum, not a slower convergence. The next section traces it to cuML's unweighted candidate reduction.

Why the array-API and cuML scalable initializers differ

Both implementations follow the same k-means|| outline. The table compares them step by step; the cuML column follows initScalableKMeansPlusPlus and kmeansPlusPlus in cuVS 26.10 cpp/src/cluster/detail/kmeans.cuh.

The candidate pools have similar sizes: 52, 290, 1700, 9328, and 46153 candidates on average from 3 to 3000 clusters. Two steps differ in substance: cuML ignores candidate weights in its greedy reduction, then repairs that with weighted Lloyd on the candidates. That recluster costs at most 0.19 s on the largest pool in this array-API reproduction.

To separate the two effects, the ablation reuses one array-API candidate pool per seed and swaps only the reduction step. The weighted Lloyd recluster mirrors cuVS's defaults and convergence test. Each variant's centers are then also refined by the benchmark's 12 farthest-relocation Lloyd iterations on the full data. Rows come from ablate_scalable_init.py on PyTorch CUDA; the cuML rows repeat the benchmark means.

The ablation explains both discrepancies:

• The no-Lloyd quality gap comes from cuML's weighted Lloyd recluster. Applying that recluster to the array-API candidates reproduces cuML's initializer objective on all five tasks, for example 1050.6 against 1091.6 at 3 clusters and 0.01957 against 0.01960 at 3000 clusters. The sampling differences, fixed-size versus Bernoulli draws and the extra uniform rows, are not what separates the two.
• The worse 3000-cluster full fit comes from the unweighted greedy step. Every unweighted variant ends near 0.0190 after Lloyd, like cuML, and every weighted variant ends near 0.0181, like greedy k-means++ on the full data. The recluster barely moves these final values. Without weights, k-means++ treats each sampled candidate as one point. A plausible mechanism, not measured here, is that D² oversampling puts many candidates in sparse regions between the 300 blobs, so too many centers start there.
• Weighted greedy followed by weighted Lloyd is the best combination here. It has the best no-Lloyd objective at 550 and 3000 clusters and the best final objective at 550 and 3000 clusters. It would cost one extra weighted Lloyd run on at most 46153 candidates.

On the 3- and 17-cluster tasks, both reclustered variants and cuML end up to 0.11% above the fits started from unrefined rows. A better starting objective does not guarantee a better 12-iteration fit.

Conclusions

• The array API with PyTorch CUDA is the fastest full fit and Lloyd-only run at 3 and 17 clusters; cuML is fastest from 100 clusters up.
• Farthest-point relocation makes all three array-API backends match scikit-learn on the only task with empty clusters. Keeping empty centers and cuML end about 7% above the best Lloyd-only result on that task.
• After 12 Lloyd iterations initialization barely matters on four tasks; everything but cuML scalable at 3000 clusters ends within 0.4% of the best.
• cuML's scalable initializer looks best before Lloyd because it runs weighted Lloyd on its candidates. It ends worse at 3000 clusters because its greedy reduction ignores the candidate weights.
• Adding a weighted Lloyd recluster to the array-API scalable initializer, while keeping its weighted greedy reduction, is a cheap candidate for a better initializer.

Previous: scalable initialization. Next: source link audit. Index: journal.