K-Means a escala: por qué falla y qué viene después

El k-means es el algoritmo que todo el mundo aprende en su curso introductorio de ML y luego usa mal el resto de su carrera. Es engañosamente simple: eliges K centros de cluster, asignas cada punto al centro más cercano, mueves los centros a la media de sus puntos asignados y repites. Tres líneas de pseudocódigo. El problema es que esa simplicidad esconde una serie de minas que explotan cuando lo aplicas a datos reales a gran escala.

Con 10.000 puntos, k-means termina en milisegundos y nadie se preocupa por los detalles de implementación. Con 10 millones de puntos, la elección del método de inicialización cambia el tiempo de ejecución por un factor de 100. Con 1.000 millones de puntos, el algoritmo estándar no cabe en memoria y necesitas enfoques fundamentalmente distintos. Un artículo reciente sobre Flash-KMeans aborda exactamente esto: logra resultados exactos de k-means con una memoria mucho menor y una convergencia más rápida. Veamos qué hace difícil a k-means a gran escala y cómo las variantes modernas lo resuelven.

Qué hace realmente el K-Means estándar

El algoritmo de Lloyd, que es lo que se entiende cuando se dice 'k-means', tiene dos pasos por iteración. El paso de asignación: para cada punto, calcula su distancia a los K centroides y lo asigna al más cercano. El paso de actualización: recalcula cada centroide como la media de todos los puntos asignados a él.

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

El paso de asignación cuesta O(n × K × d) por iteración, donde n es el número de puntos, K el número de clusters y d la dimensionalidad. Para n = 1.000 millones, K = 1.000 y d = 128 (un escenario realista de clustering de embeddings), son 128 billones de operaciones de coma flotante por iteración. Incluso con un teraflop de cómputo, son 128 segundos por iteración, y k-means normalmente necesita entre 10 y 50 iteraciones para converger.

El problema de la inicialización

Antes de que k-means ejecute una sola iteración, tiene que elegir las posiciones iniciales de los centroides. Esta elección importa mucho más de lo que la mayoría cree: una mala inicialización puede llevar a converger en una solución arbitrariamente peor que la óptima.

Inicialización aleatoria (elegir K puntos de datos al azar como centroides iniciales) es simple pero poco fiable. Si dos centroides iniciales caen en el mismo cluster, un cluster quedará sin representar. El algoritmo convergerá, pero a una solución subóptima. Con K = 100, la probabilidad de al menos una colisión es alta.

K-means++ (Arthur y Vassilvitskii, 2007) es la solución estándar. Elige el primer centroide al azar y luego cada centroide siguiente con una probabilidad proporcional a su distancia al cuadrado al centroide existente más cercano. Los puntos lejos de cualquier centroide existente tienen más probabilidad de ser elegidos, lo que reparte los centroides iniciales por los datos. Esto ofrece una garantía de aproximación O(log K): la solución está demostradamente dentro de O(log K) del óptimo.

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)

El problema: la inicialización de k-means++ en sí es O(n × K × d), porque requiere K pasadas por todo el dataset. Para K y n grandes, la inicialización tarda más que varias iteraciones del algoritmo de Lloyd. Variantes escalables como k-means|| (Bahmani et al., 2012) reducen el número de pasadas sobremuestreando candidatos en paralelo y luego consolidándolos en K centroides.

El muro de la memoria

El k-means estándar necesita tener todo el dataset en memoria a la vez. Para el paso de asignación, debes acceder a cada punto y compararlo con cada centroide. Con 1.000 millones de vectores float32 de 128 dimensiones, solo los datos ocupan 512 GB, más de lo que tienen la mayoría de las máquinas individuales.

La solución ingenua es el mini-batch k-means: en cada iteración se muestrea un subconjunto aleatorio de puntos y los centroides se actualizan con esa muestra. Funciona, pero converge más lento y hacia una solución algo distinta (normalmente algo peor) que el k-means exacto. Para muchas aplicaciones la pérdida de calidad es aceptable. Para otras, como los codebooks de cuantización en búsqueda vectorial, la diferencia sí importa.

Flash-KMeans sigue otro enfoque. En lugar de procesar todo el dataset en memoria, usa una pasada en streaming con caché inteligente. La idea clave: la mayoría de los puntos no cambian de cluster entre iteraciones. Si un punto está firmemente en el cluster 7 (mucho más cerca del centroide 7 que de cualquier otro), calcular las K distancias es trabajo desperdiciado. Manteniendo cotas de distancia, Flash-KMeans puede omitir el cálculo completo de distancias para la mayoría de los puntos en la mayoría de las iteraciones.

Optimizaciones con la desigualdad triangular

El algoritmo de Elkan (2003) usa la desigualdad triangular para omitir cálculos de distancia innecesarios. La desigualdad dice que la distancia del punto P al centroide A es como mucho la distancia de P al centroide B más la distancia de B a A. Si sabes que P está asignado actualmente al centroide B, y conoces la distancia entre los centroides A y B, a veces puedes demostrar que A está demasiado lejos de P sin calcular la distancia real.

En la práctica, esto elimina entre el 80 y el 95 % de los cálculos de distancia tras las primeras iteraciones, cuando la mayoría de los puntos ya están cerca de su centroide correcto. Los cálculos restantes son para puntos cerca de las fronteras entre clusters, los únicos que podrían cambiar realmente de asignación.

El compromiso: el algoritmo de Elkan necesita O(n × K) de memoria adicional para almacenar las cotas de distancia (cotas inferiores de cada punto a cada centroide, más cotas superiores al centroide asignado). Para K grande, este coste de memoria puede ser considerable. El algoritmo de Hamerly lo reduce a O(n) manteniendo una sola cota inferior por punto, a costa de poder omitir menos cálculos.

Aceleración con GPU

K-means es trivialmente paralelo en el paso de asignación: el cálculo de distancia de cada punto es independiente. Esto lo hace ideal para la aceleración con GPU. La biblioteca cuML de NVIDIA y FAISS de Facebook incluyen implementaciones de k-means en GPU que logran aceleraciones de 10 a 50 veces respecto a las de CPU en datasets grandes.

El matiz es que la memoria de la GPU es limitada. Una A100 tiene 80 GB, suficientes para unos 150 millones de vectores de 128 dimensiones. Los datasets más grandes requieren particionar entre varias GPUs o un enfoque en streaming, donde los datos se cargan por bloques, se procesan en la GPU y los resultados se agregan en la 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

Cuándo k-means es la elección equivocada

Antes de optimizar tu implementación de k-means, plantéate si es el algoritmo adecuado. Hace suposiciones fuertes que no se cumplen en muchos datasets reales.

• Clusters esféricos. K-means asume que los clusters son aproximadamente esféricos y de tamaño similar. Si tus clusters son alargados, tienen formas irregulares o tamaños muy distintos, k-means partirá los grandes y fusionará los pequeños. Los modelos de mezcla gaussiana (GMM) manejan clusters elípticos. DBSCAN maneja formas arbitrarias.
• K conocida. Hay que especificar el número de clusters de antemano. Si no sabes K, debes ejecutar k-means varias veces con distintos valores y usar una métrica (coeficiente de silueta, método del codo) para elegir el mejor. Esto multiplica el cómputo total por el número de valores de K que pruebas.
• Distancia euclídea. K-means usa la distancia euclídea por defecto. Para embeddings de texto, la similitud coseno suele ser más apropiada. Puedes evitarlo normalizando tus vectores con L2 (lo que hace equivalente la distancia euclídea a la coseno), pero es fácil olvidarlo.
• Sensibilidad a valores atípicos. Un solo outlier lejos de cualquier cluster arrastrará hacia él el centroide asignado. Variantes robustas como k-medoids (que usa la mediana en lugar de la media) manejan mejor los outliers, pero son más costosas.

Consejos prácticos

Después de usar k-means en datasets desde miles hasta miles de millones de puntos, esto es lo que he aprendido que más importa:

1. Usa siempre la inicialización k-means++. La inicialización aleatoria nunca merece el riesgo. La diferencia en la calidad final de los clusters suele ser del 10-30 %, y k-means++ añade una sobrecarga despreciable para K pequeñas o medianas.
2. Ejecútalo varias veces. K-means encuentra un óptimo local, no global. Ejecútalo 5-10 veces con distintas semillas aleatorias y quédate con el mejor resultado (la menor distancia total intra-cluster). Es un seguro barato contra las malas ejecuciones.
3. Normaliza tus features. Si una variable tiene rango [0, 1000000] y otra [0, 1], la de rango alto dominará el cálculo de distancias. Estandariza (resta la media y divide entre la desviación estándar) o normaliza min-max antes de agrupar.
4. Usa FAISS para clustering a gran escala. Si tienes más de un millón de puntos, el k-means de scikit-learn será lento. La implementación de FAISS está muy optimizada y soporta aceleración con GPU de serie.
5. Considera métodos aproximados para trabajo exploratorio. El mini-batch k-means es unas 10 veces más rápido que el k-means exacto y da resultados que suelen ser suficientemente buenos para explorar. Usa el k-means exacto para clustering en producción donde la calidad importe.

K-means es de esos algoritmos fáciles de usar, difíciles de usar bien y que merece la pena entender a fondo. La distancia entre una implementación ingenua y una optimizada, tanto en velocidad como en calidad del resultado, es enorme. A pequeña escala nada de esto importa. A la escala en que sí importa, entender la inicialización, la gestión de memoria, la poda de distancias y la aceleración con GPU es la diferencia entre un clustering que tarda horas y uno que tarda minutos.