Meta FAIR ускорила кластеризацию в 10–45 раз в тестах e-means на GPU

Если обучение векторного индекса занимает часы, ускорение генерации ответов эту проблему не решит. Сначала системе приходится распределить эмбеддинги по группам, чтобы затем быстрее искать похожие документы. Команда Meta FAIR предложила e-means: алгоритм обновляет эти группы небольшими порциями данных и в экспериментах заметно сокращает время кластеризации. Для разработчика практический вопрос простой: какую часть подготовки индекса можно сэкономить, сохранив качество поиска?

Почему привычный k-means заставляет ждать

Представим базу из миллионов документов. Каждый документ уже превращён в вектор — набор чисел, описывающий его содержание. Один из способов организовать поиск — разбить векторы на группы и при запросе просматривать подходящие группы. Для обучения такого разбиения часто используют k-means.

У каждой группы есть центр. Классический алгоритм Ллойда относит точки к ближайшим центрам, затем пересчитывает центры по всем назначенным точкам. Чтобы сделать одно такое обновление, приходится обработать весь обучающий набор. На больших объёмах повторение этой процедуры становится дорогим этапом подготовки данных.

e-means начинает двигать центры после небольших порций векторов. Например, очередная порция показывает, что документы о платежах расположены немного иначе, чем предполагалось сначала. Центр соответствующей группы сдвигается сразу. Старые назначения постепенно теряют вес, поэтому ранние неудачные решения меньше влияют на дальнейший расчёт.

Для каждого кластера алгоритм хранит взвешенные сумму векторов и число точек; деление одного на другое даёт центр. Но результат обеспечивается целым набором приёмов: авторы также меняют силу обновлений по ходу обучения, восстанавливают почти пустые группы и используют специальные вычислительные ядра. Размер порции и начальный шаг выбираются по формулам, полученным в экспериментах. Это избавляет от отдельного перебора этих параметров, хотя число кластеров по-прежнему задаёт пользователь. Устройство метода разобрано в описании проекта.

Что именно ускорили в экспериментах

По статье авторов, e-means достигает ошибки обычного k-means после 25 проходов за 1,4–2,6 прохода. Под ошибкой здесь понимается средний квадрат расстояния от вектора до ближайшего центра. Чем он меньше, тем лучше алгоритм решает именно задачу кластеризации.

В сравнении с оптимизированной GPU-реализацией Flash-KMeans авторы сообщают об ускорении примерно в 10–45 раз по времени. В крупном эксперименте они также разбили 10 миллиардов векторов DINO на 16 миллионов кластеров за 2,5 прохода, получив ошибку ниже результата алгоритма Ллойда после 25 проходов.

Эти числа стоит читать вместе с условиями сравнения. Соперником служит конкретный алгоритм, запущенный на определённое число итераций; выигрыш зависит от данных, оборудования и реализации. Это результаты исследовательской команды, а не обещание одинаковой экономии для любой базы. Подробности сравнений и реализации опубликованы в полном тексте работы.

Как проверить пользу для своего продукта

Для команды с регулярно обновляемым векторным индексом это повод поставить небольшой эксперимент. Возьмите репрезентативную выборку собственных эмбеддингов, сохраните одинаковое число кластеров и сравните время обучения центров. Затем постройте индексы и проверьте, какую долю нужных соседей каждый вариант находит при одинаковом бюджете поиска.

Такой второй замер принципиален: меньшая ошибка кластеризации сама по себе не гарантирует лучший ответ RAG-системы. На ответ влияют также эмбеддинги, разбиение документов, настройки поиска, переранжирование и генерация. А если основное время подготовки уходит на чтение файлов или вычисление эмбеддингов, даже большое ускорение кластеризации даст ограниченный выигрыш во всём процессе.

Начать можно с открытого кода в Faiss. Для продакта критерий успеха — сколько времени и денег экономит полный цикл обновления базы при приемлемой полноте поиска. Если кластеризация занимает существенную часть этого цикла, e-means заслуживает проверки уже сейчас.