K-Means su Larga Scala: Perché Fallisce e Cosa Viene Dopo

K-means è l'algoritmo che tutti imparano nel corso introduttivo di machine learning e poi usano male per il resto della carriera. È ingannevolmente semplice: scegli K centroidi, assegna ogni punto al centroide più vicino, sposta i centroidi nella media dei punti assegnati, ripeti. Tre righe di pseudocodice. Il problema è che questa semplicità nasconde una serie di mine che esplodono quando lo applichi a dati reali su larga scala.

Con 10.000 punti, k-means termina in millisecondi e nessuno si preoccupa dei dettagli implementativi. Con 10 milioni di punti, la scelta del metodo di inizializzazione cambia il tempo di esecuzione di 100 volte. Con 1 miliardo di punti, l'algoritmo standard non sta più in memoria e servono approcci radicalmente diversi. Un recente paper su Flash-KMeans affronta esattamente questo problema, ottenendo risultati esatti di k-means con una memoria drasticamente ridotta e una convergenza più veloce. Vediamo cosa rende difficile k-means su larga scala e come le varianti moderne lo risolvono.

Cosa Fa Davvero K-Means Standard

L'algoritmo di Lloyd, quello a cui ci si riferisce quando si dice 'k-means', ha due passi per iterazione. Il passo di assegnazione: per ogni punto, calcola la distanza da tutti i K centroidi e lo assegna al più vicino. Il passo di aggiornamento: ricalcola ogni centroide come media di tutti i punti assegnati a esso.

import numpy as np
def kmeans_lloyd(X, K, max_iter=100):
n, d = X.shape
# Random initialization (bad, but we'll fix this later)
centroids = X[np.random.choice(n, K, replace=False)]
for _ in range(max_iter):
# Assignment: O(n * K * d) — this is the bottleneck
distances = np.linalg.norm(X[:, None] - centroids[None, :], axis=2)
labels = np.argmin(distances, axis=1)
# Update: O(n * d)
new_centroids = np.array([
X[labels == k].mean(axis=0) for k in range(K)
])
if np.allclose(centroids, new_centroids):
break
centroids = new_centroids
return labels, centroids

Il passo di assegnazione costa O(n × K × d) per iterazione, dove n è il numero di punti, K il numero di cluster e d la dimensionalità. Per n = 1 miliardo, K = 1000, d = 128 (uno scenario realistico di clustering di embedding), sono 128 trilioni di operazioni in virgola mobile per iterazione. Anche con un teraflop di calcolo, fanno 128 secondi per iterazione, e k-means di solito richiede 10-50 iterazioni per convergere.

Il Problema dell'Inizializzazione

Prima che k-means esegua una sola iterazione, deve scegliere le posizioni iniziali dei centroidi. Questa scelta conta molto più di quanto la maggior parte delle persone pensi: un'inizializzazione sbagliata può portare a una convergenza su una soluzione arbitrariamente peggiore dell'ottimo.

Inizializzazione casuale (scegliere K punti dati a caso come centroidi iniziali) è semplice ma inaffidabile. Se due centroidi iniziali finiscono nello stesso cluster, un cluster resterà senza rappresentanti. L'algoritmo convergerà, ma verso una soluzione subottimale. Con K = 100, la probabilità che si verifichi almeno una collisione è alta.

K-means++ (Arthur e Vassilvitskii, 2007) è la soluzione standard. Sceglie il primo centroide a caso, poi sceglie ciascun centroide successivo con probabilità proporzionale al quadrato della sua distanza dal centroide esistente più vicino. I punti lontani da qualsiasi centroide esistente hanno più probabilità di essere scelti, distribuendo i centroidi iniziali sui dati. Questo garantisce un'approssimazione di O(log K): la soluzione è dimostrabilmente entro un fattore O(log K) dall'ottimo.

def kmeans_plus_plus_init(X, K):
n, d = X.shape
centroids = [X[np.random.randint(n)]]
for _ in range(1, K):
# Distance from each point to nearest existing centroid
dists = np.min([
np.sum((X - c) ** 2, axis=1) for c in centroids
], axis=0)
# Sample proportional to distance squared
probs = dists / dists.sum()
next_idx = np.random.choice(n, p=probs)
centroids.append(X[next_idx])
return np.array(centroids)

Il problema: l'inizializzazione k-means++ stessa costa O(n × K × d), perché richiede K passaggi sull'intero dataset. Per K e n grandi, l'inizializzazione richiede più tempo di diverse iterazioni dell'algoritmo di Lloyd. Varianti scalabili come k-means|| (Bahmani et al., 2012) riducono il numero di passaggi campionando in parallelo più candidati in eccesso, per poi consolidarli in K centroidi.

Il Muro della Memoria

K-means standard richiede che l'intero dataset sia in memoria contemporaneamente. Per il passo di assegnazione, devi accedere a ogni punto e confrontarlo con ogni centroide. Con 1 miliardo di vettori float32 a 128 dimensioni, i soli dati occupano 512 GB, più di quanto la maggior parte delle singole macchine possa contenere.

La soluzione ingenua è il mini-batch k-means: a ogni iterazione si campiona un sottoinsieme casuale di punti e si aggiornano i centroidi in base al campione. Funziona, ma converge più lentamente e verso una soluzione leggermente diversa (di solito leggermente peggiore) rispetto al k-means esatto. Per molte applicazioni la perdita di qualità è accettabile. Per altre, come i codebook di quantizzazione nella ricerca vettoriale, la differenza di qualità conta.

Flash-KMeans adotta un approccio diverso. Invece di elaborare l'intero dataset in memoria, usa un passaggio in streaming con una cache intelligente. L'intuizione chiave: la maggior parte dei punti non cambia cluster tra un'iterazione e l'altra. Se un punto è saldamente nel cluster 7 (molto più vicino al centroide 7 che a qualsiasi altro), controllare tutte le K distanze è lavoro sprecato. Mantenendo dei limiti sulle distanze, Flash-KMeans può saltare il calcolo completo della distanza per la maggior parte dei punti nella maggior parte delle iterazioni.

Ottimizzazioni con la Disuguaglianza Triangolare

L'algoritmo di Elkan (2003) usa la disuguaglianza triangolare per saltare i calcoli di distanza non necessari. La disuguaglianza dice: la distanza tra il punto P e il centroide A è al massimo la distanza tra P e il centroide B più la distanza tra B e A. Se sai che P è attualmente assegnato al centroide B e conosci la distanza tra i centroidi A e B, a volte puoi dimostrare che A è troppo lontano da P senza calcolarne la distanza effettiva.

In pratica, questo elimina l'80-95% dei calcoli di distanza dopo le prime iterazioni, quando la maggior parte dei punti è già vicina al proprio centroide corretto. I calcoli rimanenti riguardano i punti vicini ai confini dei cluster, gli unici che potrebbero effettivamente cambiare assegnazione.

Il compromesso: l'algoritmo di Elkan richiede O(n × K) di memoria aggiuntiva per memorizzare i limiti sulle distanze (limiti inferiori da ogni punto a ogni centroide, più limiti superiori al centroide assegnato). Con K grandi, questo costo di memoria può diventare sostanziale. L'algoritmo di Hamerly lo riduce a O(n) mantenendo un solo limite inferiore per punto, al costo di meno calcoli scartati.

Accelerazione con GPU

Nel passo di assegnazione, k-means è imbarazzantemente parallelo: il calcolo della distanza di ogni punto è indipendente dagli altri. Questo lo rende un candidato naturale per l'accelerazione GPU. La libreria cuML di NVIDIA e FAISS di Facebook includono entrambe implementazioni di k-means su GPU che ottengono accelerazioni di 10-50x rispetto alle implementazioni su CPU per dataset grandi.

Il problema è che la memoria della GPU è limitata. Una A100 ha 80 GB di memoria, sufficienti per circa 150 milioni di vettori a 128 dimensioni. Dataset più grandi richiedono o il partizionamento su più GPU, o un approccio in streaming in cui i dati vengono caricati a blocchi, elaborati sulla GPU e i risultati aggregati sulla CPU.

# Using FAISS for GPU-accelerated k-means
import faiss
import numpy as np
# 10 million 128-dimensional vectors
n, d, K = 10_000_000, 128, 1000
X = np.random.randn(n, d).astype('float32')
# CPU k-means (for comparison)
kmeans_cpu = faiss.Kmeans(d, K, niter=20, verbose=True)
kmeans_cpu.train(X)  # ~120 seconds
# GPU k-means (single GPU)
kmeans_gpu = faiss.Kmeans(d, K, niter=20, verbose=True, gpu=True)
kmeans_gpu.train(X)  # ~8 seconds — 15x faster

Quando K-Means È la Scelta Sbagliata

Prima di ottimizzare la tua implementazione di k-means, chiediti se è l'algoritmo giusto. Fa assunzioni forti che non valgono per molti dataset reali.

• Cluster sferici. K-means assume che i cluster siano approssimativamente sferici e di dimensioni simili. Se i tuoi cluster sono allungati, di forma irregolare o di dimensioni molto diverse, k-means dividerà quelli grandi e fonderà quelli piccoli. I Gaussian Mixture Models (GMM) gestiscono cluster ellittici. DBSCAN gestisce forme arbitrarie.
• K noto. Devi specificare il numero di cluster in anticipo. Se non conosci K, devi eseguire k-means più volte con valori diversi e usare una metrica (silhouette score, metodo del gomito) per scegliere il migliore. Questo moltiplica il costo computazionale totale per il numero di valori di K che provi.
• Distanza euclidea. K-means usa di default la distanza euclidea. Per gli embedding testuali, la similarità del coseno è di solito più appropriata. Puoi aggirare il problema normalizzando i vettori in norma L2 (il che rende la distanza euclidea equivalente alla distanza del coseno), ma è facile dimenticarsene.
• Sensibilità agli outlier. Un singolo outlier lontano da qualsiasi cluster attirerà a sé il centroide assegnato. Varianti robuste come k-medoids (che usa la mediana invece della media) gestiscono meglio gli outlier, ma sono più costose.

Consigli Pratici

Avendo usato k-means su dataset da migliaia a miliardi di punti, ecco cosa ho imparato conta di più:

1. Usa sempre l'inizializzazione k-means++. L'inizializzazione casuale non vale mai il rischio. La differenza nella qualità finale dei cluster è spesso del 10-30%, e k-means++ aggiunge un sovraccarico trascurabile per K piccoli e medi.
2. Esegui più volte. K-means trova un ottimo locale, non globale. Eseguilo 5-10 volte con seed casuali diversi e tieni il risultato migliore (la somma delle distanze intra-cluster più bassa). È un'assicurazione economica contro le esecuzioni sfortunate.
3. Normalizza le feature. Se una feature ha range [0, 1000000] e un'altra [0, 1], la feature con range più ampio dominerà il calcolo delle distanze. Standardizza (sottrai la media, dividi per la deviazione standard) o normalizza min-max prima del clustering.
4. Usa FAISS per il clustering su larga scala. Se hai più di un milione di punti, il k-means di scikit-learn sarà lento. L'implementazione di FAISS è molto ottimizzata e supporta l'accelerazione GPU già pronta all'uso.
5. Considera i metodi approssimati per l'analisi esplorativa. Il mini-batch k-means è 10 volte più veloce del k-means esatto e dà risultati di solito abbastanza buoni per esplorare i dati. Usa il k-means esatto per il clustering in produzione, dove la qualità conta.

K-means è uno di quegli algoritmi facili da usare, difficili da usare bene e che vale la pena capire a fondo. Il divario tra un'implementazione ingenua e una ottimizzata, sia in velocità che in qualità del risultato, è enorme. Su piccola scala niente di tutto questo conta. Su scala dove conta, capire inizializzazione, gestione della memoria, potatura delle distanze e accelerazione GPU fa la differenza tra un clustering che richiede ore e uno che richiede minuti.