K-Means em Escala: Por Que Falha e o Que Vem Depois

K-means é o algoritmo que todo mundo aprende na disciplina introdutória de ML e depois usa de forma incorreta pelo resto da carreira. Ele é enganosamente simples: escolha K centros de cluster, atribua cada ponto ao centro mais próximo, mova os centros para a média dos pontos atribuídos a eles e repita. Três linhas de pseudocódigo. O problema é que essa simplicidade esconde uma série de armadilhas que explodem quando você aplica o algoritmo a dados reais em escala.

Com 10.000 pontos, o k-means termina em milissegundos e ninguém se preocupa com detalhes de implementação. Com 10 milhões de pontos, a escolha do método de inicialização muda o tempo de execução em 100 vezes. Com 1 bilhão de pontos, o algoritmo padrão simplesmente não cabe na memória e você precisa de abordagens fundamentalmente diferentes. Um artigo recente sobre o Flash-KMeans ataca exatamente esse problema, alcançando resultados exatos de k-means com uso de memória bem menor e convergência mais rápida. Vamos entender o que torna o k-means difícil em escala e como as variantes modernas resolvem isso.

O Que o K-Means Padrão Realmente Faz

O algoritmo de Lloyd, que é o que as pessoas querem dizer quando falam em 'k-means', tem duas etapas por iteração. Na etapa de atribuição, para cada ponto de dados, calculamos a distância até todos os K centroides e o atribuímos ao mais próximo. Na etapa de atualização, recalculamos cada centroide como a média de todos os pontos atribuídos a ele.

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

A etapa de atribuição custa O(n × K × d) por iteração, onde n é o número de pontos, K é o número de clusters e d é a dimensionalidade. Para n = 1 bilhão, K = 1000 e d = 128 (um cenário realista de clustering de embeddings), isso dá 128 trilhões de operações de ponto flutuante por iteração. Mesmo com um teraflop de poder de computação, são 128 segundos por iteração, e o k-means normalmente precisa de 10 a 50 iterações para convergir.

O Problema da Inicialização

Antes de o k-means executar qualquer iteração, ele precisa escolher as posições iniciais dos centroides. Essa escolha importa muito mais do que a maioria das pessoas percebe: uma inicialização ruim pode levar à convergência para uma solução arbitrariamente pior que a ótima.

Inicialização aleatória (escolher K pontos de dados aleatórios como centroides iniciais) é simples, mas pouco confiável. Se dois centroides iniciais caírem no mesmo cluster, um cluster ficará sem representação. O algoritmo vai convergir, mas para uma solução subótima. Com K = 100, a probabilidade de pelo menos uma colisão é alta.

K-means++ (Arthur e Vassilvitskii, 2007) é a solução padrão para isso. Ele escolhe o primeiro centroide aleatoriamente e depois escolhe cada centroide seguinte com probabilidade proporcional à sua distância quadrática até o centroide existente mais próximo. Pontos distantes de qualquer centroide existente têm mais chance de serem escolhidos, espalhando os centroides iniciais pelos dados. Isso oferece uma garantia de aproximação O(log K): a solução está comprovadamente dentro de um fator O(log K) do ótimo.

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)

O porém: a própria inicialização do k-means++ é O(n × K × d), porque exige K passagens pelo conjunto de dados inteiro. Para K e n grandes, a inicialização leva mais tempo do que várias iterações do algoritmo de Lloyd. Variantes escaláveis como o k-means|| (Bahmani et al., 2012) reduzem o número de passagens ao fazer sobreamostragem de candidatos em paralelo e, depois, consolidá-los em K centroides.

A Parede de Memória

O k-means padrão precisa ter o conjunto de dados inteiro na memória ao mesmo tempo. Na etapa de atribuição, você precisa acessar cada ponto e compará-lo com cada centroide. Com 1 bilhão de vetores de 128 dimensões em float32, só os dados ocupam 512 GB, mais do que a maioria das máquinas individuais tem.

A solução ingênua é o mini-batch k-means: a cada iteração, amostramos um subconjunto aleatório de pontos e atualizamos os centroides com base nessa amostra. Funciona, mas converge mais devagar e para uma solução um pouco diferente (geralmente um pouco pior) do que o k-means exato. Para muitas aplicações, a perda de qualidade é aceitável. Para outras, como codebooks de quantização em busca vetorial, a diferença de qualidade importa.

O Flash-KMeans adota outra abordagem. Em vez de processar o conjunto de dados inteiro na memória, ele usa uma passagem em streaming com cache inteligente. A ideia central é que a maioria dos pontos não muda de cluster entre as iterações. Se um ponto está firmemente no cluster 7 (muito mais perto do centroide 7 do que de qualquer outro centroide), calcular todas as K distâncias é trabalho desperdiçado. Mantendo limites de distância, o Flash-KMeans consegue pular o cálculo completo de distância para a maioria dos pontos na maioria das iterações.

Otimizações com a Desigualdade Triangular

O algoritmo de Elkan (2003) usa a desigualdade triangular para pular cálculos de distância desnecessários. Ela diz que a distância do ponto P até o centroide A é no máximo a distância de P até o centroide B somada à distância de B até A. Se você sabe que P está atualmente atribuído ao centroide B e conhece a distância entre os centroides A e B, às vezes consegue provar que A está longe demais de P sem calcular a distância real.

Na prática, isso elimina de 80% a 95% dos cálculos de distância depois das primeiras iterações, quando a maioria dos pontos já está perto do centroide correto. Os cálculos restantes são para pontos próximos às fronteiras dos clusters, que são os únicos que realmente podem mudar de atribuição.

O porém: o algoritmo de Elkan exige O(n × K) de memória adicional para guardar os limites de distância (limites inferiores de cada ponto até cada centroide, mais limites superiores até o centroide atribuído). Para K grande, esse custo de memória pode ser substancial. O algoritmo de Hamerly reduz isso para O(n) ao manter apenas um limite inferior por ponto, em troca de menos cálculos evitados.

Aceleração com GPU

O k-means é naturalmente paralelo na etapa de atribuição: o cálculo de distância de cada ponto é independente. Isso o torna um candidato natural para aceleração em GPU. A biblioteca cuML da NVIDIA e o FAISS do Facebook incluem implementações de k-means em GPU que alcançam ganhos de 10 a 50 vezes em relação às implementações em CPU para conjuntos de dados grandes.

A complicação é que a memória da GPU é limitada. Uma A100 tem 80 GB de memória, o suficiente para cerca de 150 milhões de vetores de 128 dimensões. Conjuntos maiores exigem particionamento em múltiplas GPUs ou uma abordagem em streaming, em que os dados são carregados em blocos, processados na GPU e os resultados agregados na 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 o K-Means É a Escolha Errada

Antes de otimizar sua implementação de k-means, considere se ele é o algoritmo certo. Ele faz suposições fortes que não valem para muitos conjuntos de dados reais.

• Clusters esféricos. O K-means assume que os clusters são aproximadamente esféricos e de tamanhos parecidos. Se seus clusters são alongados, têm formatos irregulares ou tamanhos muito diferentes, o k-means vai dividir os grandes e unir os pequenos. Modelos de Mistura Gaussiana (GMMs) lidam com clusters elípticos. DBSCAN lida com formatos arbitrários.
• K conhecido. Você precisa especificar o número de clusters de antemão. Se não sabe qual é o K, precisa rodar o k-means várias vezes com valores diferentes e usar uma métrica (coeficiente de silhueta, método do cotovelo) para escolher o melhor. Isso multiplica o custo computacional total pelo número de valores de K testados.
• Distância euclidiana. O K-means usa distância euclidiana por padrão. Para embeddings de texto, a similaridade de cosseno costuma ser mais apropriada. Você pode contornar isso normalizando seus vetores com L2 (o que torna a distância euclidiana equivalente à distância de cosseno), mas é fácil esquecer disso.
• Sensibilidade a outliers. Um único outlier distante de qualquer cluster vai puxar o centroide atribuído até ele. Variantes robustas como o k-medoids (que usa a mediana em vez da média) lidam melhor com outliers, mas são mais caras.

Conselhos Práticos

Depois de usar k-means em conjuntos de dados que vão de milhares a bilhões de pontos, estes são os aprendizados que mais importam:

1. Sempre use a inicialização k-means++. A inicialização aleatória nunca vale o risco. A diferença na qualidade final dos clusters costuma ser de 10% a 30%, e o k-means++ adiciona pouco overhead para K pequeno ou médio.
2. Rode várias vezes. O K-means encontra um ótimo local, não global. Rode de 5 a 10 vezes com sementes aleatórias diferentes e fique com o melhor resultado (menor distância total intra-cluster). É um seguro barato contra execuções ruins.
3. Normalize suas features. Se uma feature vai de [0, 1000000] e outra de [0, 1], a feature de faixa maior vai dominar o cálculo de distância. Padronize (subtraia a média e divida pelo desvio padrão) ou normalize por min-max antes de fazer o clustering.
4. Use o FAISS para clustering em larga escala. Se você tem mais de um milhão de pontos, o k-means do scikit-learn vai ficar lento. A implementação do FAISS é altamente otimizada e suporta aceleração por GPU sem configuração extra.
5. Considere métodos aproximados para trabalho exploratório. O mini-batch k-means é 10 vezes mais rápido que o k-means exato e entrega resultados geralmente bons o suficiente para exploração. Use o k-means exato em clustering de produção, onde a qualidade importa.

O K-means é um desses algoritmos fáceis de usar, difíceis de usar bem e que vale a pena entender a fundo. A distância entre uma implementação ingênua e uma otimizada, tanto em velocidade quanto em qualidade do resultado, é enorme. Em pequena escala, nada disso importa. Na escala em que importa, entender inicialização, gerenciamento de memória, poda de distâncias e aceleração por GPU é a diferença entre um clustering que leva horas e um que leva minutos.