K-means: When Clustering Becomes a Feature Engineering Trick

Lectures 8a, 8b, and 8c, and the course changes gears for the first time: every model up to now learned to predict a label already sitting in the data (y). K-means gets no label at all, it just finds groups.

The problem: separating without knowing the answer

The professor goes back to Iris, but this time hides the species (y) and only passes the two petal variables (length and width) to the algorithm. The question shifts from "which species is this flower?" to "how many natural groups exist here, and where do they sit?"

K-means by hand

from sklearn.base import BaseEstimator, ClusterMixin, TransformerMixin

class KMeans(BaseEstimator, ClusterMixin, TransformerMixin):
    def __init__(self, n_clusters=3, max_iter=100):
        self.n_clusters = n_clusters
        self.max_iter = max_iter

    def fit(self, X, y=None):
        self.centroids = X[random.sample(range(len(X)), self.n_clusters)]
        max_iter = self.max_iter
        while max_iter > 0:
            max_iter -= 1
            y_pred = self.predict(X)
            for i in range(self.n_clusters):
                self.centroids[i] = np.mean(X[y_pred==i], axis=0)
            if np.allclose(self.centroids, self.previous_centroids[-1], atol=1e-9):
                break
        return self

    def predict(self, X):
        # assigns every point to its nearest centroid
        ...

Notice the class signature: ClusterMixin and TransformerMixin together, the first time a class in this playlist inherits from two mixins at once on top of BaseEstimator. That's not an accident, and the reason becomes clear further down the post.

The algorithm (called Lloyd's K-means, the most common one) is just two steps repeated until nothing changes:

1. Assign: every point goes to its nearest centroid (Euclidean distance).
2. Recompute: every centroid becomes the mean of the points assigned to it.

Bishop formalizes both steps as minimizing a distortion measure,

$J = \sum_{n=1}^{N}\sum_{k=1}^{K} r_{nk}\|\mathbf{x}_n - \boldsymbol{\mu}_k\|^2$

where $r_{nk}=1$ if point $n$ belongs to group $k$ (and 0 otherwise). Fixing the centroids, the optimal $r_{nk}$ is obvious (assign to the nearest one, exactly step 1). Fixing the $r_{nk}$, the optimal centroid is the mean of the points assigned to it, because that's where $J$'s derivative hits zero (exactly step 2). Every step can only decrease $J$ or leave it unchanged, never increase it, so the algorithm always converges, it just might land on a local minimum, not necessarily the best possible grouping (which is why, in practice, scikit-learn runs K-means several times with different starting centroids and keeps the result with the lowest $J$).

Interactive: watching the centroids move

My own reconstruction of the same algorithm (assign, recompute, repeat), on the 150 real Iris points, the same two variables from the decision trees post. Click "Step" and watch the inertia ($J$) drop each round, and the regions (who belongs to which centroid) settle:

The points stay colored by true species, just so you can visually compare: the regions K-means finds on its own, never seeing the label, match the real species pretty well, especially for Setosa (already clearly separated back in the decision trees post). Click "Reset" a few times: depending on where the centroids land by chance, the final result can shift a bit, exactly the local-minimum limitation Bishop describes.

Picking K: the elbow method

One question was left open: how do you choose how many groups (K) to look for? Scikit-learn's K-means exposes inertia_, which is exactly Bishop's $J$ at the end of training:

kmeans_k3 = KMeans(n_clusters=3).fit(X)
kmeans_k8 = KMeans(n_clusters=8).fit(X)

Output: inertia with K=3: 31.37. With K=8: 8.31.

More groups always reduces (or ties) inertia, in the limit (K = number of points) it hits zero, every point becomes its own group, which is useless. The trick is to plot inertia against several K values and look for the elbow: the point where the curve stops dropping fast and starts dropping slowly. Before the elbow, each extra group still captures real structure. After it, each extra group is just slicing up noise. The professor repeats this plot for Iris, for the wine dataset that already showed up in this playlist, and for Digits (1797 images of handwritten digits, 8×8 pixels each), without pulling a fixed number out of any of them, the point is just to show the curve's shape.

The trick: using the groups as a new feature

This is where lecture 8b flips the whole post on its head. Instead of using K-means only to cluster, the professor uses it inside a supervised pipeline, as a preprocessing step:

from sklearn.linear_model import RidgeClassifier
from sklearn.model_selection import cross_val_score
from sklearn.pipeline import make_pipeline

model = RidgeClassifier()
scores = cross_val_score(model, X, y, cv=5)

Output (Digits, RidgeClassifier alone, straight on the 64 pixels): 0.888 average accuracy.

model = make_pipeline(
    KMeans(n_clusters=50),
    RidgeClassifier()
)
scores = cross_val_score(model, X, y, cv=5)

Output (same dataset, with a 50-cluster K-means ahead of the classifier): 0.939.

A serious jump, just from slotting a K-means in the middle. How does this work? This is where the TransformerMixin I mentioned earlier comes back into play. Inside a Pipeline, every step except the last needs .transform(), not .predict(). Scikit-learn's K-means implements both: .predict() returns the index of the nearest group (a single number), but .transform() returns the distance to each of the K centroids (K numbers). A 64-pixel image becomes a vector of 50 numbers, each one saying "how similar is this image to the prototype of group 1, group 2, ..., group 50." Every centroid works as a "prototype digit" learned with no label at all, and the distance to each one becomes a new feature, more informative to the linear classifier than a raw pixel.

The professor still uses Optuna to search for the ideal number of groups (between 10 and 200):

Output: the best found was n_clusters=136, with 0.963 average cross-validation accuracy. Trying n_clusters=250 by hand (more groups than Optuna even tried): 0.966, even better.

More groups, more different "prototypes" to compare against, more information for the final classifier, at least as far as the professor tested.

The same trick, on a bigger dataset: MNIST

aula08c repeats the exact recipe on MNIST (60 thousand training images, 28×28 pixels each, so 784 variables per image, quite a bit more than Digits):

Eight percentage points of gain, on the same dataset, on the same linear classifier, just by swapping "64/784 raw pixel values" for "250 distances to prototypes learned with no label." The idea of reusing an unsupervised algorithm as a feature source for a supervised problem is more general than it looks: any "distance to a prototype" vector carries information the raw pixel doesn't carry on its own.

Wrapping up

Practical application

I reproduced the same recipe (K-means as a transformer ahead of RidgeClassifier) on Digits, with a fixed seed and n_init=10 (K-means runs 10 times with different starting centroids and keeps the one with the lowest inertia, scikit-learn's default), to confirm the effect reproducibly.

cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
ridge_scores = cross_val_score(RidgeClassifier(random_state=42), X, y, cv=cv)
for k in [50, 136, 250]:
    model = make_pipeline(KMeans(n_clusters=k, random_state=42, n_init=10), RidgeClassifier(random_state=42))
    scores = cross_val_score(model, X, y, cv=cv)

My numbers come out higher than the original notebook's (0.888/0.939/0.963), likely because I fixed n_init=10 (K-means tries several initializations and keeps the best, avoiding the bad local minima Bishop warned about above) and a seed my environment could reproduce stably. But the trend is identical to the notebook's: more groups, better accuracy, and the improvement from "raw pixel" to "distance to prototype" is large and consistent at every group count tested.