e-means : k-means en 1 à 3 époques au lieu de 25

Ce qui se passe

Sept auteurs (Swann Bessa, Pierre-Emmanuel Mazaré, Gergely Szilvasy, Hugo Touvron, David Picard, Matthijs Douze et Hervé Jégou), dont des chercheurs de Meta FAIR, proposent e-means, une remplaçante de l'algorithme de Lloyd, le k-means classique. Le texte est une prépublication (arXiv:2611.09501), non relue par des pairs : les chiffres ci-dessous sont ceux des auteurs, sur leurs bancs.

Pourquoi ça compte

Le k-means sert à entraîner les centroïdes d'un index vectoriel ou d'une quantification. Lloyd recalcule les centroïdes sur tout le jeu de données à chaque itération, en général 20 à 50 fois. Sur des dizaines de millions de vecteurs, c'est le coût dominant. Selon les auteurs, e-means atteint une erreur de test inférieure à celle de Lloyd à 25 époques en 1,2 à 3 époques, soit :

10 à 45 fois moins de temps de calcul sur GPU, face à Flash-KMeans, l'implémentation Lloyd de référence ;
8 à 15 fois sur CPU, face à Faiss.

Sur les 11 jeux d'embeddings du banc VIBE (texte, code, image, multimodal ; dimensions 200 à 3 072 ; 8 192 ou 16 384 clusters), il faut entre 1,6 et 2,4 époques. Un exemple chiffré : sur un H100, 10 millions de vecteurs et 65 536 clusters, 25 époques de Lloyd en PyTorch pur prennent 142 à 155 secondes, et 14 à 22 secondes avec Flash-KMeans.

Le mécanisme

e-means travaille sur de petits lots. Pour chaque cluster, il garde une somme et un effectif, et les met à jour par une moyenne mobile exponentielle de pas α : les anciennes affectations s'oublient au rythme α, alors que le k-means par lots classique cumule tout l'historique. Le centroïde est le rapport somme sur effectif. Lloyd n'en est qu'un cas extrême : lot égal au jeu de données entier, α = 1.

Quatre ajouts évitent le réglage à la main :

la taille de lot et le pas sont déduits de lois d'échelle à partir de n (points), k (clusters) et T (époques) ;
une initialisation revue, et un pas qui décroît en cosinus jusqu'à 0,1 fois sa valeur de départ ;
la réanimation des clusters morts : un cluster trop petit est refait en coupant un gros cluster en deux ;
une étape finale qui termine sur la moyenne exacte de chaque cluster à la dernière époque.

Les auteurs expliquent aussi le gain par une meilleure sortie des minima locaux : sur un jeu jouet de deux amas de tailles très différentes (k = 8), Lloyd reste piégé dans 26 % des 200 tirages, e-means dans 3,5 %.

Comment s'y prendre

Le code vit dans Faiss, sous faiss/contrib/e_means :

Python

from faiss.contrib.e_means import Emeans   # NumpyEmeans pour NumPy

em = Emeans(d, k, nepochs=T)               # T époques, 5 par défaut
centroids, labels = em.train(x)            # x : tenseur PyTorch (n, d)

labels donne le cluster de chaque vecteur à la dernière époque. Comparez sur vos embeddings : même k, même jeu de test, et regardez à la fois l'erreur quadratique moyenne et le temps. La validation des auteurs porte sur 24 jeux de données et configurations, dont 234 réglages comparés à k-means sur 16 jeux de 100 000 à 100 millions de points.

Source : page du projet e-means, en anglais, avec l'article arXiv:2611.09501 et le dépôt faiss/contrib/e_means.