Onderzoekers van Meta FAIR hebben e-means gepubliceerd, een algoritme dat wordt gepresenteerd als een directe vervanging van het Lloyd’s k-means algoritme. Volgens de auteurs bereikt het in 1,4 tot 2,6 epochs het verlies dat Lloyd’s k-means in 25 epochs behaalt. Het is beschikbaar in FAISS 1.16.0, onder de module faiss.contrib.e_means.
K-means is een clusteringsmethode die gegevens verdeelt in groepen, elk vertegenwoordigd door een centrum genaamd centroïde. Voor wie vectorindexen, quantizers of datavoorbereidingspipelines bouwt, is deze berekening vaak de traagste offline stap, en dat is precies wat e-means beoogt te versnellen.
Een voortschrijdend gemiddelde in plaats van een volledige herberekening
Alle varianten van k-means houden voor elke groep een som S en een teller N bij, waarbij de centroïde het quotiënt van S door N is. Ze verschillen vooral in hoe S en N evolueren. Het Lloyd-algoritme herberekent deze waarden vanuit de volledige dataset bij elke iteratie. MacQueen voegt in 1967 continu de statistieken van batches toe, zodat oude toewijzingen nooit verdwijnen.
E-means kiest een tussenweg: het dempt oude statistieken met een snelheid α, volgens de formule S ← (1−α)S + αs, en hetzelfde voor N. Oude toewijzingen vervagen zo geleidelijk, en de centroïden volgen de recente gegevens. Het geval van Lloyd komt overeen met een volledige batch met α = 1. Op de DEEP-dataset bevindt dit geval zich 1,8% boven het beste behaalde verlies.
De auteurs hebben ook een klassieke gradient descent getest. In zijn beste configuratie blijft deze 1,6% boven de EMA op DEEP en 0,4% op SIFT. Zij vermoeden dat de EMA-methode wint omdat de centroïden altijd binnen de convexe omhulling van hun toegewezen punten blijven.
Parameters vastgelegd door schaalwetten
Twee empirisch afgestelde schaalwetten bepalen de batchgrootte B en de initiële stap α0 op basis van de grootte van de dataset, het aantal clusters en het aantal epochs. Op datasets die buiten beschouwing werden gehouden, bereikt de mediane spijt 0,07% voor de stapwet. De stap volgt vervolgens een cosinusschema, van α0 tot 0,1α0. Twee andere mechanismen completeren het systeem: het herstarten van ‘dode’ clusters, d.w.z. die onder een drempelwaarde voor grootte vallen en worden herstart door een groot cluster te splitsen, en een afronding op het exacte gemiddelde van de laatste epoch.
Volgens de auteurs zijn de aangekondigde prestaties 10 tot 45 keer sneller op GPU dan Flash-KMeans, en 8 tot 15 keer sneller op CPU dan Faiss k-means. Op de VIBE-benchmark hadden de 11 embedding-datasets tussen 1,6 en 2,4 epochs nodig. Voor schaalvergroting werd de DINO10B-dataset, bestaande uit 10 miljard embeddings van 1024 dimensies, verdeeld over 16,7 miljoen clusters in 2,5 epochs. De GPU-resultaten werden behaald op H100’s. Voor configuraties met 16 miljoen centroïden bereiken de accumulatoren in float32 64 GiB, wat vereist dat de centroïden over meerdere GPU’s worden verdeeld.
Een demonstratie op een enkele CPU-thread, gepubliceerd in de aankondigingsdraad, geeft 4,39 s voor Lloyd over 25 epochs, tegen 0,43 s voor e-means over 2,2 epochs. E-means wordt gedistribueerd onder MIT-licentie, via FAISS.
