Des chercheurs de Meta FAIR ont publié e-means, un algorithme présenté comme un remplacement direct de l’algorithme k-means de Lloyd. Selon les auteurs, il atteint en 1,4 à 2,6 époques la perte que le k-means de Lloyd obtient en 25 époques. Il est disponible dans FAISS 1.16.0, sous le module faiss.contrib.e_means.
- Meta FAIR publie e-means, un remplaçant direct de k-means de Lloyd, disponible dans FAISS 1.16.0.
- e-means atteint en 1,4 à 2,6 époques la perte que Lloyd obtient en 25 époques.
- Il lisse les statistiques anciennes par une moyenne mobile exponentielle, avec des paramètres fixés par des lois d’échelle.
- Les auteurs annoncent des gains de 10 à 45 fois sur GPU et de 8 à 15 fois sur CPU, sous licence MIT.
Le k-means est une méthode de regroupement qui répartit des données en groupes, chacun représenté par un centre appelé centroïde. Pour ceux qui construisent des index vectoriels, des quantificateurs ou des pipelines de préparation de données, ce calcul est souvent l’étape hors ligne la plus lente, et c’est précisément celle que vise e-means.
Une moyenne mobile plutôt qu’un recalcul complet
Toutes les variantes du k-means conservent, pour chaque groupe, une somme S et un compteur N, le centroïde étant le quotient de S par N. Elles diffèrent surtout par la manière dont S et N évoluent. L’algorithme de Lloyd recalcule ces valeurs à partir de l’ensemble des données à chaque itération. MacQueen, en 1967, ajoute en continu les statistiques des lots, de sorte que les anciennes affectations ne disparaissent jamais.
e-means adopte une voie intermédiaire : il atténue les anciennes statistiques à un taux α, selon la formule S ← (1−α)S + αs, et de même pour N. Les affectations anciennes s’effacent ainsi peu à peu, et les centroïdes suivent les données récentes. Le cas de Lloyd correspond à un lot complet avec α = 1. Sur le jeu de données DEEP, ce cas se situe 1,8 % au-dessus de la meilleure perte obtenue.
Les auteurs ont aussi testé une descente de gradient classique. Dans sa meilleure configuration, elle reste à 1,6 % au-dessus de l’EMA sur DEEP et à 0,4 % sur SIFT. Ils supposent que la méthode EMA l’emporte parce que les centroïdes restent toujours à l’intérieur de l’enveloppe convexe de leurs points affectés.
Des paramètres fixés par des lois d’échelle
Deux lois d’échelle, ajustées empiriquement, déterminent la taille de lot B et le pas initial α0 à partir de la taille du jeu de données, du nombre de clusters et du nombre d’époques. Sur des jeux de données tenus à l’écart, la médiane du regret atteint 0,07 % pour la loi du pas. Le pas suit ensuite un calendrier cosinus, de α0 jusqu’à 0,1α0. Deux autres mécanismes complètent le dispositif : la relance des clusters « morts », c’est-à-dire ceux qui tombent sous un seuil de taille et sont relancés en scindant un grand cluster, et une finalisation sur la moyenne exacte de la dernière époque.
Selon les auteurs, les performances annoncées sont de 10 à 45 fois plus rapides sur GPU que Flash-KMeans, et de 8 à 15 fois plus rapides sur CPU que Faiss k-means. Sur le benchmark VIBE, les 11 jeux de données d’embeddings ont nécessité entre 1,6 et 2,4 époques. Pour le passage à l’échelle, le jeu DINO10B, composé de 10 milliards d’embeddings de 1 024 dimensions, a été réparti en 16,7 millions de clusters en 2,5 époques. Les résultats sur GPU ont été obtenus sur des H100. Pour les configurations à 16 millions de centroïdes, les accumulateurs en float32 atteignent 64 Gio, ce qui impose de répartir les centroïdes sur plusieurs GPU.
Une démonstration sur un seul thread CPU, publiée dans le fil d’annonce, donne 4,39 s pour Lloyd sur 25 époques, contre 0,43 s pour e-means sur 2,2 époques. e-means est distribué sous licence MIT, via FAISS.
