Investigadores da Meta FAIR publicaram e-means, um algoritmo apresentado como um substituto direto do algoritmo k-means de Lloyd. Segundo os autores, ele atinge em 1,4 a 2,6 épocas a perda que o k-means de Lloyd obtém em 25 épocas. Está disponível no FAISS 1.16.0, sob o módulo faiss.contrib.e_means.
O k-means é um método de agrupamento que distribui dados em grupos, cada um representado por um centro chamado centroide. Para quem constrói índices vetoriais, quantificadores ou pipelines de preparação de dados, este cálculo é frequentemente a etapa offline mais lenta, e é precisamente essa que o e-means visa.
Uma média móvel em vez de um recálculo completo
Todas as variantes do k-means mantêm, para cada grupo, uma soma S e um contador N, sendo o centroide o quociente de S por N. Elas diferem sobretudo pela forma como S e N evoluem. O algoritmo de Lloyd recalcula estes valores a partir do conjunto de dados em cada iteração. MacQueen, em 1967, adiciona continuamente as estatísticas dos lotes, de modo que as antigas atribuições nunca desaparecem.
O e-means adota uma via intermédia: atenua as estatísticas antigas a uma taxa α, segundo a fórmula S ← (1−α)S + αs, e o mesmo para N. As atribuições antigas apagam-se assim pouco a pouco, e os centroides seguem os dados recentes. O caso de Lloyd corresponde a um lote completo com α = 1. No conjunto de dados DEEP, este caso situa-se 1,8 % acima da melhor perda obtida.
Os autores também testaram uma descida de gradiente clássica. Na sua melhor configuração, ela mantém-se 1,6 % acima da EMA no DEEP e 0,4 % no SIFT. Supõem que o método EMA vence porque os centroides permanecem sempre dentro do invólucro convexo dos seus pontos atribuídos.
Parâmetros fixados por leis de escala
Duas leis de escala, ajustadas empiricamente, determinam o tamanho do lote B e o passo inicial α0 a partir do tamanho do conjunto de dados, do número de clusters e do número de épocas. Em conjuntos de dados mantidos à parte, a mediana do arrependimento atinge 0,07 % para a lei do passo. O passo segue depois um calendário cosseno, de α0 até 0,1α0. Dois outros mecanismos completam o dispositivo: o relançamento dos clusters «mortos», ou seja, aqueles que caem abaixo de um limiar de tamanho e são relançados dividindo um grande cluster, e uma finalização sobre a média exata da última época.
Segundo os autores, os desempenhos anunciados são de 10 a 45 vezes mais rápidos em GPU que o Flash-KMeans, e de 8 a 15 vezes mais rápidos em CPU que o Faiss k-means. No benchmark VIBE, os 11 conjuntos de dados de embeddings necessitaram entre 1,6 e 2,4 épocas. Para a passagem à escala, o conjunto DINO10B, composto por 10 mil milhões de embeddings de 1024 dimensões, foi distribuído em 16,7 milhões de clusters em 2,5 épocas. Os resultados em GPU foram obtidos em H100. Para configurações com 16 milhões de centroides, os acumuladores em float32 atingem 64 GiB, o que obriga a distribuir os centroides por várias GPUs.
Uma demonstração numa única thread de CPU, publicada no fio de anúncio, dá 4,39 s para Lloyd em 25 épocas, contra 0,43 s para e-means em 2,2 épocas. O e-means é distribuído sob licença MIT, via FAISS.
