-
Notifications
You must be signed in to change notification settings - Fork 662
[BUG] KMeans ignores sample_weight when computing inertia and score #8530
Description
Describe the bug
cuml.cluster.KMeans appears to ignore sample_weight when computing both the fitted inertia_ and the value returned by score().
With non-uniform training weights, cuML returns the unweighted inertia:
inertia: 8.0
instead of the weighted value 12.0. When scoring two samples with weights [2, 2], cuML similarly returns the unweighted score -4.0 instead of -8.0.
The equivalent scikit-learn call applies the supplied weights to both calculations.
Steps/Code to reproduce bug
cuML reproducer:
import numpy as np from cuml.cluster import KMeans model = KMeans(n_clusters=2,init=np.array([[1., 1.],[11., 11.],]),n_init=1,).fit(np.array([[0., 0.],[1., 1.],[2., 2.],[10., 10.],[11., 11.],[12., 12.],]), sample_weight=np.array([1., 1., 1., 2., 2., 2.])) print("inertia:", model.inertia_) print("score:",model.score(np.array([[0., 0.],[10., 10.],]),sample_weight=np.array([2., 2.])))
Output:
inertia: 8.0
score: -4.0
For comparison, the equivalent scikit-learn code:
import numpy as np from sklearn.cluster import KMeans model = KMeans(n_clusters=2,init=np.array([[1., 1.],[11., 11.],]),n_init=1,).fit(np.array([[0., 0.],[1., 1.],[2., 2.],[10., 10.],[11., 11.],[12., 12.],]), sample_weight=np.array([1., 1., 1., 2., 2., 2.])) print("inertia:", model.inertia_) print("score:",model.score(np.array([[0., 0.],[10., 10.],]),sample_weight=np.array([2., 2.])))
Output:
inertia: 12.0
score: -8.0
Expected behavior
sample_weight should scale each sample's contribution to the K-means objective.
For the training data in this reproducer, the squared distances to the assigned centers are [2, 0, 2] in each cluster. With training weights [1, 1, 1, 2, 2, 2], the expected inertia is:
(1 * 2 + 1 * 0 + 1 * 2) + (2 * 2 + 2 * 0 + 2 * 2) = 12
For the two scoring samples, both squared distances are 2; weights [2, 2] therefore give an objective of 8 and an expected score of -8.
Environment details (please complete the following information):
- Environment location: Docker
- Linux Distro/Architecture: Ubuntu 24.04 / x86_64
- GPU Model/Driver: NVIDIA GeForce RTX 4090 / 595.71.05
- CUDA: 13.2
- Method of cuDF & cuML install: conda
conda list:
conda list
# packages in environment at /opt/conda/envs/rapids-26.08:
#
# Name Version Build Channel
# Name Version Build Channel
python 3.14.6 h242f9ac_102_cp314 conda-forge
numpy 2.4.6 py314h2b28147_0 conda-forge
scipy 1.16.3 py314hf07bd8e_2 conda-forge
scikit-learn 1.9.0 np2py314hf09ca88_0 conda-forge
rapids 26.08.00 cuda13_260806_c2656556 rapidsai
cuml 26.08.00 cuda13_cp311_abi3_260805_265b9da6 rapidsai
libcuml 26.08.00 cuda13_260805_265b9da6 rapidsai
cudf 26.08.00 cuda13_cp311_abi3_260805_ff5b362d rapidsai
libraft 26.08.00 cuda13_260805_ebf92684 rapidsai
libraft-headers 26.08.00 cuda13_260805_ebf92684 rapidsai
pylibraft 26.08.00 cuda13_cp311_abi3_260805_ebf92684 rapidsai
cuvs 26.08.01 cuda13_cp311_abi3_260806_25b1be43 rapidsai
libcuvs 26.08.01 cuda13_260806_25b1be43 rapidsai
cupy 14.1.1 py314hdea9c46_0 conda-forge
cupy-core 14.1.1 py314hcd3b49b_0 conda-forge
numba 0.64.0 py314h8169c2f_0 conda-forge
numba-cuda 0.30.4 py314h42812f9_0 conda-forge
rmm 26.08.00 cuda13_cp311_abi3_260805_42d059f1 rapidsai
librmm 26.08.00 cuda13_260805_42d059f1 rapidsai
cuda-version 13.3 hcbadf70_3 conda-forge
cuda-bindings 13.3.1 py314h42812f9_1 conda-forge
cuda-cudart 13.3.29 hecca717_0 conda-forge
cuda-nvrtc 13.3.33 hecca717_0 conda-forge
libcublas 13.6.0.2 h676940d_0 conda-forge
libcusolver 12.2.6.9 h676940d_0 conda-forge
libcusparse 12.8.2.51 hecca717_0 conda-forge
libcurand 10.4.3.29 h676940d_0 conda-forge
Additional context
The initial centers are supplied explicitly and n_init=1, making the result deterministic. The per-cluster weights are constant, so the fitted centers remain [1, 1] and [11, 11] whether weighted or unweighted. This isolates the issue to the weighted objective calculation rather than center movement.
The cuML results exactly match the unweighted calculations:
unweighted training inertia = 2 +たす 0 +たす 2 +たす 2 +たす 0 +たす 2 =わ 8
unweighted scoring objective = 2 + 2 = 4
The cuML KMeans documentation documents sample_weight for both fit() and score() as the weights for each observation.