Суть прорыва: от полных батчей к экспоненциальному сглаживанию
Классический алгоритм k-means (Lloyd) требует пересчета сумм и счетчиков по всему датасету на каждой итерации, что делает его узким местом при работе с миллиардами векторов. Meta FAIR предложила замену — e-means (exponential moving average k-means). Алгоритм использует мини-батчи и обновляет центроиды через экспоненциальное скользящее среднее (EMA), позволяя «старым» назначениям затухать. Это не только ускоряет сходимость, но и помогает избегать локальных минимумов, в которых часто застревает классический Lloyd.
Ключевые метрики производительности
Алгоритм протестирован на данных от 100K до 10 млрд точек с количеством кластеров до 16.7 млн. По заявленным данным авторов, e-means превосходит базовые решения по скорости и эффективности использования ресурсов:
| Метрика | Результат e-means | Сравнение / Примечание |
|---|---|---|
| Эпохи для сходимости | 1.4 – 2.6 | Достигает потерь Lloyd’s k-means (25 эпох) в ~10 раз быстрее |
| GPU ускорение | 10x – 45x | Против Flash-KMeans при равном уровне потерь (H100) |
| CPU ускорение | 8x – 15x | Против Faiss k-means |
| Масштаб | 10 млрд точек | Разбиение DINO10B (1024 dim) на 16.7M кластеров за 2.5 эпохи |
| Отказ от тюнинга | Автоматический | Размеры батча и шаг определяются масштабными законами, а не подбором |
Технические детали и оптимизации
Для минимизации ручного тюнинга команда внедрила масштабные законы, которые автоматически вычисляют размер батча (B) и начальный шаг (α0) на основе размера датасета (n), количества кластеров (k) и целевых эпох (T). Медианная регрета на отложенных данных составила всего 0.07%.
Также реализованы критически важные для больших данных механизмы:
- Rebirth (Возрождение кластеров): пустые кластеры (мертвые узлы) реанимируются путем разделения крупных кластеров. В тесте на 16,384 кластеров это снизило долю мертвых кластеров с 26.3% до 0.3%.
- Fused GPU kernels: операция поиска ближайших центроидов объединена в один ядро, что экономит секунды на каждом этапе.
- Finalization: обучение завершается вычислением точного среднего значения для каждого кластера на последней эпохе.
Почему это важно для разработчиков
Кластеризация часто является самым медленным этапом офлайн-обработки данных при создании векторных индексов, квантизаторов или пайплайнов курирования данных. e-means, доступный в FAISS 1.16.0 через faiss.contrib.e_means, позволяет сократить время подготовки данных в десятки раз. Алгоритм работает как на CPU, так и на GPU (с поддержкой PyTorch тензоров и NumPy), что делает его универсальным инструментом для современных AI-инфраструктур.
Источник: MarkTechPost ↗