k-Means Clustering
Before diving into K-Means, understand the paradigm shift:
Unsupervised Learning — Quick Foundation
Before diving into K-Means, understand the paradigm shift:
| Aspect | Supervised | Unsupervised |
|---|---|---|
| Input | $(X, y)$ — features + labels | $X$ only — features |
| Goal | Predict $y$ from $X$ | Find structure in $X$ |
| Examples | Classification, regression | Clustering, dim reduction, density estimation |
| Evaluation | Accuracy, MSE on labeled test set | Harder — no ground truth |
| Algorithms | LR, GBDT, NNs | K-Means, DBSCAN, PCA, GMM |
Common unsupervised tasks:
- Clustering — group similar samples (K-Means, DBSCAN, hierarchical).
- Dimensionality reduction — compress to fewer features (PCA, t-SNE, UMAP).
- Density estimation — model $P(x)$ (GMM, KDE).
- Anomaly detection — find unusual samples (Isolation Forest, One-Class SVM).
- Association rule mining — find patterns (Apriori).
Warning: Why unsupervised learning is hard
- No ground truth — how do you know if a clustering is “right”?
- Evaluation is subjective — depends on downstream use.
- Hyperparameter selection is trickier — no labels to validate against.
What is K-Means?
K-Means is a partitional clustering algorithm that divides $n$ samples into $K$ non-overlapping clusters. Each sample belongs to exactly one cluster — the one whose centroid (mean) is closest.
Core idea:
- Each cluster is represented by its centroid — the mean of all points in the cluster.
- Each point is assigned to the nearest centroid (Euclidean distance).
- Centroids are updated to be the means of their assigned points.
- This process iterates until convergence.
“K-Means” comes from:
- $K$: number of clusters (a hyperparameter).
- Means: each cluster is represented by the mean of its points.
The K-Means Algorithm — Step by Step
Input: Dataset $X = {x_1, x_2, \dots, x_n}$, number of clusters $K$. Output: $K$ centroids and cluster assignments.
Step 1 — Initialize: choose $K$ initial centroids $\mu_1, \mu_2, \dots, \mu_K$.
Step 2 — Assignment: assign each point to the nearest centroid:
\[c_i = \arg\min_{k} \|x_i - \mu_k\|^2\]Step 3 — Update: recompute centroids as the mean of assigned points:
\[\mu_k = \frac{1}{|C_k|}\sum_{i \in C_k} x_i\]Step 4 — Repeat steps 2–3 until convergence (centroids stop moving, or max iterations reached).
Note: Convergence
- The algorithm is guaranteed to converge — each step decreases or maintains the objective.
- However, it may converge to a local minimum, not the global one.
- Convergence usually happens in 10–30 iterations for typical data.
The Objective Function — What K-Means Optimizes
K-Means minimizes the within-cluster sum of squares (WCSS), also called inertia:
\[J = \sum_{k=1}^K \sum_{x_i \in C_k} \|x_i - \mu_k\|^2\]- Lower $J$ → tighter, more compact clusters.
- $J = 0$ means every point is a centroid (only when $K = n$).
EM Perspective
K-Means is a special case of Expectation-Maximization (EM):
- E-step: assign points to clusters (computes “expected” labels given current centroids).
- M-step: update centroids to maximize fit (compute means given assignments).
Each step decreases $J$, so the algorithm converges. But to a local minimum, since the objective is non-convex in cluster assignments.
Mathematical Note
The update step (centroid = mean of points) is provably optimal because the mean minimizes squared distances:
\[\mu^* = \arg\min_\mu \sum_i \|x_i - \mu\|^2 \Rightarrow \mu^* = \frac{1}{n}\sum_i x_i\]This is why K-Means uses Euclidean distance and means specifically — they’re mathematically coupled.
Initialization — A Critical Step
The initial centroids significantly affect the final clustering. Bad initialization → bad clusters.
Random Initialization (Forgy Method)
Randomly pick $K$ data points as initial centroids. Simple, but can yield poor results and is highly sensitive to seed.
K-Means++ (The Standard Today)
A smarter initialization scheme (Arthur & Vassilvitskii, 2007) that drastically improves clustering quality.
Algorithm:
- Pick the first centroid uniformly at random from the data.
- For each subsequent centroid:
- Compute $D(x)$ = distance from each point to the nearest already-chosen centroid.
- Pick the next centroid with probability proportional to $D(x)^2$.
- Continue until $K$ centroids are chosen.
- Run standard K-Means from these centroids.
Check: Why K-Means++ works
- Spreads initial centroids far apart — closer to the global optimum.
- Theoretical guarantee: expected error is $O(\log K)$ times the optimal.
- In practice, dramatically faster convergence and better solutions.
- Used in scikit-learn by default (
init='k-means++').
Multiple Initializations
Run K-Means multiple times with different initializations and pick the result with lowest inertia. In scikit-learn: n_init=10 (default) means try 10 different starts.
Choosing K — The Big Question
$K$ is a hyperparameter that must be specified upfront — but how do you choose it?
The Elbow Method
Plot inertia vs $K$. Look for an “elbow” where the rate of decrease sharply slows.
Adding more clusters always decreases inertia, but at some point the gain becomes marginal. The elbow is where you get diminishing returns.
Warning: Issue: often subjective — sometimes there’s no clear elbow.
Silhouette Score
For each point $i$:
\[s_i = \frac{b_i - a_i}{\max(a_i, b_i)}\]Where:
- $a_i$ = average distance from $i$ to other points in its own cluster (cohesion).
- $b_i$ = average distance from $i$ to points in the nearest other cluster (separation).
Range: $[-1, 1]$
| Score | Interpretation |
|---|---|
| Near +1 | Well-clustered |
| Near 0 | On the boundary between clusters |
| Near −1 | Probably in the wrong cluster |
Pick $K$ with the highest average silhouette score.
Other Methods
- Gap Statistic — compare inertia to a random uniform reference distribution; pick $K$ where the gap is largest.
- Davies-Bouldin Index — ratio of within-cluster scatter to between-cluster separation; lower is better.
- Domain knowledge — often the most practical; you know how many clusters make sense (e.g., 5 customer segments).
Practical Workflow
- Try several $K$ values (e.g., 2–10).
- Plot elbow + silhouette.
- Look for consensus.
- Apply domain knowledge.
Note: Sometimes there’s no single “right” $K$ — just useful ones.
Distance Metrics in K-Means
K-Means is defined for Euclidean (L2) distance — the centroid (mean) is provably optimal under squared Euclidean distance. Using other distances breaks the mathematical guarantee that the mean minimizes the objective.
Alternatives for Other Distances
| Algorithm | Center type | Distance | Notes |
|---|---|---|---|
| K-Medoids (PAM) | Actual data points (medoids) | Any | More flexible, slower |
| K-Medians | Median of points | Manhattan | Robust to outliers |
| Spherical K-Means | Normalized centroids | Cosine | Standard for text/embeddings |
Tip: Cosine similarity trick If using K-Means with cosine similarity, normalize features to unit length first — then Euclidean distance on unit vectors is equivalent to cosine.
Feature Scaling — Critical!
K-Means is distance-based, so feature scale matters enormously.
Example without scaling:
- Feature 1: age (0–100)
- Feature 2: income (0–1,000,000)
Income completely dominates the distance computation — clusters form based on income alone.
Always standardize features before K-Means:
- Z-score: $x’ = \frac{x - \mu}{\sigma}$ (mean 0, std 1).
- Min-max: $x’ = \frac{x - x_{\min}}{x_{\max} - x_{\min}}$ (range $[0, 1]$).
Forgetting to scale features is one of the most common mistakes with K-Means. Always scale first.
K-Means Assumptions (and When They Break)
| Assumption | What breaks it |
|---|---|
| Clusters are spherical and isotropic | Elongated, banana-shaped, or irregular clusters |
| Clusters are roughly the same size | Imbalanced cluster sizes — gets split or merged incorrectly |
| Clusters have similar density | Dense cluster next to a sparse one — misallocation |
| Data is continuous and numerical | Categorical features (need encoding or K-Modes) |
| Euclidean distance is meaningful | High-dim, sparse, or non-numerical data |
| $K$ is known in advance | Requires choosing $K$ — often the hardest part |
When K-Means Fails
- Non-convex clusters (e.g., two interlocked half-moons).
- Clusters of very different sizes or densities.
- High-dimensional data (curse of dimensionality).
- Noisy data with many outliers.
Tip: Alternatives when K-Means fails
- DBSCAN — density-based; handles arbitrary shapes and outliers.
- Spectral Clustering — graph-based; handles non-convex clusters.
- Hierarchical Clustering — no need to specify $K$ upfront.
- Gaussian Mixture Models (GMM) — soft probabilistic cluster assignments; handles elliptical clusters.
K-Means Stopping Criteria — Deep Dive
K-Means is iterative — it keeps alternating between assignment and update steps. Without a stopping rule:
- The algorithm could run forever (in theory).
- Or waste compute on tiny improvements that don’t matter.
- Or stop too early before reaching a good solution.
A good stopping criterion balances quality, efficiency, and robustness. In practice, K-Means usually converges in 10–30 iterations for typical data, so stopping criteria mostly serve as safeguards rather than strict optimality controls.
The Main Stopping Criteria
Production implementations like scikit-learn typically use multiple criteria simultaneously — stopping as soon as any is satisfied.
(A) Centroids Stop Moving (Convergence)
The classic criterion: stop when centroids stabilize between iterations:
\[\max_{k} \|\mu_k^{(t)} - \mu_k^{(t-1)}\| < \text{tol}\]Or with sum of squared shifts:
\[\sum_{k=1}^K \|\mu_k^{(t)} - \mu_k^{(t-1)}\|^2 < \text{tol}\]If centroids barely moved this iteration, further iterations won’t change much.
- Threshold
tol: small positive number (e.g., $10^{-4}$). - In scikit-learn: the
tolparameter (default1e-4); checks Frobenius norm of centroid shift.
(B) Cluster Assignments Stop Changing
Stop when no points change clusters between iterations.
If every point stays in the same cluster, centroids will compute the same means → algorithm has converged exactly.
- This is the strictest convergence criterion — guarantees fixed-point convergence.
- In practice, always satisfied when (A) is satisfied with small enough tolerance.
- Often combined with (A) for safety.
Warning: Edge case: ties With ties (a point equidistant from multiple centroids), assignments may oscillate. Tie-breaking rules (e.g., pick lowest cluster ID) prevent infinite loops.
(C) Inertia Stops Decreasing
Stop when the objective function (WCSS) decrease becomes negligible:
\[\frac{J^{(t-1)} - J^{(t)}}{J^{(t-1)}} < \text{tol}\]Or absolute change:
\[|J^{(t-1)} - J^{(t)}| < \text{tol}\]- Since K-Means monotonically decreases inertia, the change is always non-negative.
- Easy to compute — already calculated each iteration.
- Provides a smooth, principled stopping signal.
(D) Maximum Iterations
A hard cap on iterations regardless of convergence.
- Safety net — guarantees the algorithm terminates.
- Resource control — caps compute cost.
- Production reliability — prevents pathological cases from hanging.
Typical default: 100–300 iterations (scikit-learn default: max_iter=300).
Warning: If
max_iteris hit, something is wrong — bad initialization, weird data — and you should investigate.
(E) Wall-Clock Time Limit (Production)
In some production settings, cap by time rather than iterations (e.g., “run for at most 30 seconds”). Less common but pragmatic when integrated into latency-sensitive pipelines.
How They’re Combined in Practice
Real implementations run multiple criteria simultaneously and stop as soon as any is satisfied:
while iteration < max_iter:
new_centroids = update_step(...)
if shift(new_centroids, old_centroids) < tol: # Criterion A
break
if assignments_unchanged(): # Criterion B
break
old_centroids = new_centroids
iteration += 1
# Hit max_iter (Criterion D) → exit loop
Scikit-learn specifically:
- Uses Frobenius norm of centroid shift
< tol(criterion A) as primary. - Also checks if inertia stops decreasing (criterion C).
- Has hard
max_itercap (criterion D).
Why K-Means Is Guaranteed to Converge
K-Means optimizes:
\[J = \sum_{k=1}^K \sum_{x_i \in C_k} \|x_i - \mu_k\|^2\]Each step decreases or maintains $J$:
- Assignment step: each point goes to its nearest centroid. Moving a point closer always reduces $J$.
- Update step: setting $\mu_k$ = mean of $C_k$ is provably the value that minimizes $\sum_{i \in C_k}|x_i - \mu_k|^2$ for fixed $C_k$. So the update can only decrease or maintain $J$.
Therefore: $J$ is monotonically non-increasing.
Why it must terminate:
- $J \geq 0$ — sum of squared distances is non-negative.
- $J$ is non-increasing each iteration.
- There are finitely many possible cluster assignments ($K^n$ to be precise — finite for any finite dataset).
- If $J$ ever doesn’t strictly decrease, the assignments must repeat — and once they repeat, the centroids repeat, and the algorithm has converged.
Note: Conclusion K-Means must terminate in a finite number of iterations, even without explicit stopping criteria. The stopping criteria exist for efficiency and robustness, not for theoretical convergence guarantees.