k-means clustering

Deep dive

How it works

In: n points of dimension d, the cluster count k, a seed · Out: k centroids and a label per point · State: the centroids between iterations · Cost: O(n·k·d) per iteration, usually 10-100 iterations

1. Seed (k-means++). First centre uniformly at random. Keep D(x) = squared distance to the nearest chosen centre. Pick the next centre with probability D(x)/ΣD. Update D. Repeat to k. (Greedy variant: sample 2 + ln k candidates per round and keep the one that lowers ΣD most; scikit-learn does this.)
2. Assign. a(i) = argmin_j ‖xᵢ − c_j‖², computed as ‖c_j‖² − 2xᵢ·c_j (the ‖xᵢ‖² term does not affect argmin).
3. Update. c_j = mean of points with a(i) = j.
4. Empty cluster. Re-seed it at the point with the largest current distance, or split the biggest cluster.
5. Stop when the relative inertia change or centroid movement is below tol, or after max_iter.
6. Restarts. Run several seeds and keep the lowest inertia if quality matters.

iteration:   assign (GEMM + argmin)  →  update (scatter-add sums, divide by counts)
             n×d · d×k = n×k scores      k×d sums, k counts

Code

pub struct Rng(pub u64);
impl Rng {
    pub fn f64(&mut self) -> f64 {
        self.0 ^= self.0 << 13; self.0 ^= self.0 >> 7; self.0 ^= self.0 << 17;
        (self.0 >> 11) as f64 / (1u64 << 53) as f64
    }
}

fn dist2(a: &[f32], b: &[f32]) -> f32 { a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum() }

/// k-means++ seeding (D² sampling). x: n×d row-major.
pub fn kmeanspp(x: &[f32], n: usize, d: usize, k: usize, rng: &mut Rng) -> Vec<f32> {
    let mut c = Vec::with_capacity(k * d);
    let first = (rng.f64() * n as f64) as usize;
    c.extend_from_slice(&x[first * d..(first + 1) * d]);
    let mut dmin: Vec<f32> = (0..n).map(|i| dist2(&x[i * d..(i + 1) * d], &c[..d])).collect();
    for _ in 1..k {
        let total: f64 = dmin.iter().map(|&v| v as f64).sum();
        let mut target = rng.f64() * total;
        let mut pick = n - 1;
        for (i, &v) in dmin.iter().enumerate() { target -= v as f64; if target <= 0.0 { pick = i; break; } }
        let new = x[pick * d..(pick + 1) * d].to_vec();
        for i in 0..n { dmin[i] = dmin[i].min(dist2(&x[i * d..(i + 1) * d], &new)); }
        c.extend_from_slice(&new);
    }
    c
}

/// Lloyd iterations. Returns (centroids, assignment, inertia, iterations).
pub fn lloyd(x: &[f32], n: usize, d: usize, mut c: Vec<f32>, k: usize, max_iter: usize, tol: f64)
             -> (Vec<f32>, Vec<u32>, f64, usize) {
    let mut assign = vec![0u32; n];
    let mut prev = f64::INFINITY;
    let xn: Vec<f32> = (0..n).map(|i| x[i * d..(i + 1) * d].iter().map(|v| v * v).sum()).collect();
    for it in 1..=max_iter {
        let cn: Vec<f32> = (0..k).map(|j| c[j * d..(j + 1) * d].iter().map(|v| v * v).sum()).collect();
        let mut inertia = 0.0f64;
        let mut far = (0usize, -1.0f32);
        for i in 0..n {                                   // in production: one GEMM, then row argmin
            let xi = &x[i * d..(i + 1) * d];
            let (mut best, mut bj) = (f32::INFINITY, 0);
            for j in 0..k {
                let dot: f32 = xi.iter().zip(&c[j * d..(j + 1) * d]).map(|(a, b)| a * b).sum();
                let s = cn[j] - 2.0 * dot;                // ‖x‖² is constant per row
                if s < best { best = s; bj = j; }
            }
            let dd = (best + xn[i]).max(0.0);             // clamp tiny negatives from cancellation
            if dd > far.1 { far = (i, dd); }
            inertia += dd as f64;
            assign[i] = bj as u32;
        }
        let mut sum = vec![0.0f64; k * d];
        let mut cnt = vec![0usize; k];
        for i in 0..n {
            let j = assign[i] as usize;
            cnt[j] += 1;
            for t in 0..d { sum[j * d + t] += x[i * d + t] as f64; }
        }
        for j in 0..k {
            if cnt[j] == 0 {                              // empty cluster: re-seed at the farthest point
                c[j * d..(j + 1) * d].copy_from_slice(&x[far.0 * d..(far.0 + 1) * d]);
                far.1 = -1.0;
            } else {
                for t in 0..d { c[j * d + t] = (sum[j * d + t] / cnt[j] as f64) as f32; }
            }
        }
        if (prev - inertia) / inertia.max(1e-30) < tol { return (c, assign, inertia, it); }
        prev = inertia;
    }
    (c, assign, prev, max_iter)
}

fn main() {
    let (n, d, k) = (3000, 2, 3);
    let mut rng = Rng(7);
    let centres = [(-5.0, 0.0), (0.0, 5.0), (5.0, 0.0)];
    let mut x = Vec::with_capacity(n * d);
    for i in 0..n {
        let (cx, cy) = centres[i % 3];
        let (u1, u2) = (1.0 - rng.f64(), rng.f64());      // Box-Muller normal noise
        let r = (-2.0 * u1.ln()).sqrt();
        x.push((cx + r * (6.283185307 * u2).cos()) as f32);
        x.push((cy + r * (6.283185307 * u2).sin()) as f32);
    }
    let seeds = kmeanspp(&x, n, d, k, &mut rng);
    let (c, _, inertia, iters) = lloyd(&x, n, d, seeds, k, 300, 1e-4);
    println!("converged in {iters} iterations, inertia/n = {:.3} (≈ 2 for unit-variance 2-D blobs)", inertia / n as f64);
    for j in 0..k { println!("centroid {j}: ({:.2}, {:.2})", c[j * d], c[j * d + 1]); }
}

Complexity and constants

• Assignment: 2·n·k·d FLOPs as a GEMM. For n = 10⁶, k = 4096, d = 128: 10¹² FLOPs, order of 10 ms on a modern GPU, order of seconds on a multicore CPU with BLAS.
• Update: O(n·d) scatter-add, memory-bound.
• Memory: the n×k score matrix is never needed in full. Process rows in blocks and keep only the running argmin.
• Iterations: typically 10-100 to reach tol = 1e-4. Worst-case Lloyd can take exponentially many, but that does not happen on real data.
• Training a 65536-centroid IVF quantizer on 10⁷ sampled vectors is a routine GPU job (Faiss).

Efficiency

• Cost: assignment is 2·n·k·d FLOPs as a GEMM: n = 10⁶, k = 4096, d = 128 is 10¹² FLOPs, order of 10 ms on a modern GPU, seconds on a multicore CPU (this page). Centroids take k·d·4 bytes (2 MiB here, in cache); the update is an O(n·d) memory-bound scatter-add. Assignment is compute-bound, update is bandwidth-bound.
• Levers:
  • Distances as ‖x‖² − 2x·c + ‖c‖² with a GEMM — BLAS or tensor-core speed (this page; Faiss) — large k and d.
  • Train on a sample — Faiss subsamples to 256 points per centroid by default and warns below 39 (Faiss) — IVF and codebook training.
  • Triangle-inequality bounds — skip most distance computations at identical results (Elkan 2003; Hamerly 2010) — low to moderate d; gain shrinks at high d.
  • Mini-batch updates — orders of magnitude less computation than batch k-means (Sculley 2010) — huge n, slightly worse objective.
• Naive → standard → state of the art:
  • Per-point loop over centroids, random init · scalar O(n·k·d), many iterations
  • GEMM-based Lloyd with k-means++ · 10 ms per iteration at the numbers above (GPU, order of)
  • GPU Faiss on a sample, or mini-batch · 65,536 centroids on 10⁷ vectors is a routine job (this page)

Implementation notes

• Subsample for training when n is huge: Faiss caps at 256 points per centroid by default and warns below 39.
• Use the GEMM formulation in fp32 with a tiled matmul (tiled matmul), then argmin per row. For bf16 inputs, keep the norm terms in fp32: the cancellation in ‖x‖² − 2x·c + ‖c‖² loses precision when points are far from the origin. Centre the data first.
• Elkan or Hamerly bounds (triangle inequality) skip most distance computations in late iterations on CPU when d is small to moderate.
• Mini-batch k-means: per batch, assign, then move each centre toward its batch points with per-centre step 1/count_j. Much faster for large n, slightly worse inertia.
• Parallel: shard rows across threads or machines; each computes partial sums and counts; reduce; divide. One all-reduce of k·(d+1) values per iteration.
• Scale features first. k-means is distance-based, and a feature in large units dominates.

Pitfalls

• Random-uniform seeding: two seeds in one true cluster, one cluster split in half, a local optimum Lloyd cannot leave. Use k-means++.
• Choosing k by the inertia minimum: inertia decreases monotonically with k. Use the elbow, silhouette, or a downstream metric (recall for IVF).
• Cosine similarity on unnormalized vectors with Euclidean k-means. L2-normalize the inputs (and optionally the centroids every iteration: spherical k-means).
• Empty clusters left as NaN centroids after division by zero.
• Single run on hard data: inertia varies between seeds. Use several restarts when quality matters.

Variants and alternatives

In production

Faiss (faiss.Kmeans, GPU, IVF and PQ codebook training), scikit-learn KMeans (Lloyd and Elkan) and MiniBatchKMeans, Spark MLlib (k-means||), ScaNN and other vector-search libraries for partition training, image colour quantization.

Implementations

Sources

• Lloyd, Least Squares Quantization in PCM, IEEE Transactions on Information Theory 1982. paper
• Arthur and Vassilvitskii, k-means++: The Advantages of Careful Seeding, SODA 2007. paper
• Bahmani et al., Scalable K-Means++, VLDB 2012. paper
• Sculley, Web-Scale K-Means Clustering, WWW 2010. paper
• Elkan, Using the Triangle Inequality to Accelerate k-Means, ICML 2003. paper
• Hamerly, Making k-means Even Faster, SDM 2010. paper
• Johnson, Douze, Jégou, Billion-Scale Similarity Search with GPUs, IEEE Transactions on Big Data 2019. paper · code