Naukowcy z Meta FAIR opublikowali e-means, algorytm przedstawiany jako bezpośredni zamiennik algorytmu k-means Lloyda. Według autorów osiąga on w 1,4 do 2,6 epoki stratę, którą k-means Lloyda uzyskuje w 25 epokach. Jest dostępny w FAISS 1.16.0, w module faiss.contrib.e_means.
K-means to metoda grupowania, która dzieli dane na grupy, z których każdą reprezentuje środek zwany centroidem. Dla osób budujących indeksy wektorowe, kwantyzatory lub potoki przygotowania danych, to obliczenie jest często najwolniejszym etapem offline i właśnie na niego celuje e-means.
Średnia krocząca zamiast pełnego przeliczenia
Wszystkie warianty k-means przechowują dla każdej grupy sumę S i licznik N, gdzie centroid jest ilorazem S przez N. Różnią się one głównie sposobem ewolucji S i N. Algorytm Lloyda przelicza te wartości od nowa na podstawie całego zbioru danych w każdej iteracji. MacQueen w 1967 roku stale dodaje statystyki z partii, tak że stare przypisania nigdy nie znikają.
e-means wybiera drogę pośrednią: wygładza stare statystyki współczynnikiem α, według wzoru S ← (1−α)S + αs, i analogicznie dla N. Stare przypisania stopniowo zanikają, a centroidy podążają za świeżymi danymi. Przypadek Lloyda odpowiada pełnej partii z α = 1. Na zbiorze danych DEEP ten przypadek znajduje się 1,8% powyżej najlepszej osiągniętej straty.
Autorzy przetestowali również klasyczny spadek gradientu. W swojej najlepszej konfiguracji pozostaje on na poziomie 1,6% powyżej EMA na DEEP i 0,4% na SIFT. Zakładają, że metoda EMA wygrywa, ponieważ centroidy zawsze pozostają wewnątrz otoczki wypukłej swoich przypisanych punktów.
Parametry ustalone przez prawa skalowania
Dwa prawa skalowania, dopasowane empirycznie, określają rozmiar partii B i krok początkowy α0 na podstawie rozmiaru zbioru danych, liczby klastrów i liczby epok. Na zbiorach danych odłożonych na bok mediana żalu osiąga 0,07% dla prawa kroku. Krok następnie podlega harmonogramowi cosinusowemu, od α0 do 0,1α0. Dwa inne mechanizmy uzupełniają system: restartowanie « martwych » klastrów, czyli tych, które spadają poniżej progu wielkości i są restartowane przez podzielenie dużego klastra, oraz finalizacja na dokładnej średniej z ostatniej epoki.
Według autorów, ogłoszone wydajności są od 10 do 45 razy szybsze na GPU niż Flash-KMeans i od 8 do 15 razy szybsze na CPU niż Faiss k-means. Na benchmarku VIBE, 11 zbiorów danych embeddingów wymagało od 1,6 do 2,4 epoki. W przypadku skalowania, zbiór DINO10B, składający się z 10 miliardów embeddingów o 1024 wymiarach, został podzielony na 16,7 miliona klastrów w 2,5 epoki. Wyniki na GPU uzyskano na H100. Dla konfiguracji z 16 milionami centroidów, akumulatory float32 osiągają 64 GiB, co wymaga rozłożenia centroidów na wielu GPU.
Demonstracja na pojedynczym wątku CPU, opublikowana w wątku ogłoszeniowym, pokazuje 4,39 s dla Lloyda na 25 epokach, wobec 0,43 s dla e-means na 2,2 epokach. e-means jest dystrybuowany na licencji MIT, poprzez FAISS.
