I ricercatori di Meta FAIR hanno pubblicato e-means, un algoritmo presentato come sostituto diretto dell’algoritmo k-means di Lloyd. Secondo gli autori, raggiunge in 1,4-2,6 epoche la perdita che il k-means di Lloyd ottiene in 25 epoche. È disponibile in FAISS 1.16.0, sotto il modulo faiss.contrib.e_means.
Il k-means è un metodo di clustering che suddivide i dati in gruppi, ciascuno rappresentato da un centro chiamato centroide. Per chi costruisce indici vettoriali, quantizzatori o pipeline di preparazione dei dati, questo calcolo è spesso la fase offline più lenta, ed è proprio quella che e-means mira a risolvere.
Una media mobile invece di un ricalcolo completo
Tutte le varianti del k-means conservano, per ogni gruppo, una somma S e un contatore N, dove il centroide è il quoziente di S diviso N. Differiscono principalmente nel modo in cui S e N evolvono. L’algoritmo di Lloyd ricalcola questi valori dall’intero dataset a ogni iterazione. MacQueen, nel 1967, aggiunge continuamente le statistiche dei batch, in modo che le assegnazioni precedenti non scompaiano mai.
e-means adotta una via intermedia: attenua le vecchie statistiche a un tasso α, secondo la formula S ← (1−α)S + αs, e allo stesso modo per N. Le assegnazioni vecchie si cancellano così gradualmente e i centroidi seguono i dati recenti. Il caso di Lloyd corrisponde a un batch completo con α = 1. Sul dataset DEEP, questo caso si trova all’1,8% sopra la migliore perdita ottenuta.
Gli autori hanno anche testato una classica discesa del gradiente. Nella sua configurazione migliore, rimane all’1,6% sopra l’EMA su DEEP e allo 0,4% su SIFT. Suppongono che il metodo EMA prevalga perché i centroidi rimangono sempre all’interno dell’inviluppo convesso dei loro punti assegnati.
Parametri fissati da leggi di scala
Due leggi di scala, aggiustate empiricamente, determinano la dimensione del batch B e il passo iniziale α0 a partire dalla dimensione del dataset, dal numero di cluster e dal numero di epoche. Su dataset tenuti da parte, la mediana del rimpianto raggiunge lo 0,07% per la legge del passo. Il passo segue poi una schedulazione coseno, da α0 fino a 0,1α0. Due altri meccanismi completano il dispositivo: il riavvio dei cluster « morti », cioè quelli che scendono sotto una soglia di dimensione e vengono riavviati dividendo un cluster grande, e una finalizzazione sulla media esatta dell’ultima epoca.
Secondo gli autori, le prestazioni annunciate sono 10-45 volte più veloci su GPU rispetto a Flash-KMeans, e 8-15 volte più veloci su CPU rispetto a Faiss k-means. Sul benchmark VIBE, gli 11 dataset di embedding hanno richiesto tra 1,6 e 2,4 epoche. Per lo scaling, il dataset DINO10B, composto da 10 miliardi di embedding a 1024 dimensioni, è stato suddiviso in 16,7 milioni di cluster in 2,5 epoche. I risultati su GPU sono stati ottenuti su H100. Per configurazioni con 16 milioni di centroidi, gli accumulatori in float32 raggiungono 64 GiB, il che impone di distribuire i centroidi su più GPU.
Una dimostrazione su un singolo thread CPU, pubblicata nel thread di annuncio, dà 4,39 s per Lloyd su 25 epoche, contro 0,43 s per e-means su 2,2 epoche. e-means è distribuito con licenza MIT, tramite FAISS.
