Исследователи Meta FAIR опубликовали e-means, алгоритм, представленный как прямая замена алгоритма k-means Ллойда. По словам авторов, он достигает за 1,4–2,6 эпох того уровня потерь, который k-means Ллойда достигает за 25 эпох. Он доступен в FAISS 1.16.0 в модуле faiss.contrib.e_means.
K-means — это метод кластеризации, который распределяет данные по группам, каждая из которых представлена центром, называемым центроидом. Для тех, кто строит векторные индексы, квантизаторы или конвейеры подготовки данных, это вычисление часто является самым медленным офлайн-этапом, и именно на него нацелен e-means.
Скользящее среднее вместо полного пересчета
Все варианты k-means сохраняют для каждого кластера сумму S и счетчик N, где центроид является частным от деления S на N. Они различаются в основном тем, как изменяются S и N. Алгоритм Ллойда пересчитывает эти значения на основе всего набора данных на каждой итерации. Маккуин в 1967 году непрерывно добавляет статистику батчей, так что старые назначения никогда не исчезают.
E-means выбирает промежуточный путь: он ослабляет старую статистику со скоростью α по формуле S ← (1−α)S + αs, и аналогично для N. Таким образом, старые назначения постепенно стираются, а центроиды следуют за новыми данными. Случай Ллойда соответствует полному батчу с α = 1. На наборе данных DEEP этот случай находится на 1,8 % выше наилучшего достигнутого уровня потерь.
Авторы также протестировали классический градиентный спуск. В своей лучшей конфигурации он остается на 1,6 % выше EMA на DEEP и на 0,4 % на SIFT. Они предполагают, что метод EMA побеждает потому, что центроиды всегда остаются внутри выпуклой оболочки своих назначенных точек.
Параметры, заданные законами масштабирования
Два закона масштабирования, подобранные эмпирически, определяют размер батча B и начальный шаг α0 на основе размера набора данных, количества кластеров и числа эпох. На отложенных наборах данных медианное значение сожаления достигает 0,07 % для закона шага. Затем шаг следует косинусному расписанию от α0 до 0,1α0. Два других механизма дополняют систему: перезапуск «мертвых» кластеров, то есть тех, которые падают ниже порога размера и перезапускаются путем разделения большого кластера, и финализация на точном среднем значении последней эпохи.
По словам авторов, заявленная производительность составляет от 10 до 45 раз выше на GPU, чем у Flash-KMeans, и от 8 до 15 раз выше на CPU, чем у Faiss k-means. В бенчмарке VIBE 11 наборов данных эмбеддингов потребовали от 1,6 до 2,4 эпох. Что касается масштабирования, набор DINO10B, состоящий из 10 миллиардов эмбеддингов размерностью 1024, был распределен по 16,7 миллионам кластеров за 2,5 эпохи. Результаты на GPU были получены на H100. Для конфигураций с 16 миллионами центроидов аккумуляторы в float32 достигают 64 ГБ, что требует распределения центроидов по нескольким GPU.
Демонстрация на одном потоке CPU, опубликованная в анонсе, показывает 4,39 с для Ллойда за 25 эпох против 0,43 с для e-means за 2,2 эпохи. E-means распространяется под лицензией MIT через FAISS.
