K-Means в масштабе: почему он ломается и что делать

K-means — это алгоритм, который все проходят на вводном курсе по ML, а потом всю карьеру используют неправильно. Он обманчиво прост: выбираем K центров кластеров, относим каждую точку к ближайшему центру, сдвигаем центры в среднее своих точек и повторяем. Три строки псевдокода. Но эта простота прячет целую россыпь мин, которые взрываются, когда вы применяете его к реальным данным в большом масштабе.

На 10 000 точек k-means отрабатывает за миллисекунды, и никто не думает о деталях реализации. На 10 миллионах выбор метода инициализации меняет время работы в 100 раз. На миллиарде точек стандартный алгоритм просто не помещается в память, и нужны принципиально другие подходы. Недавняя статья про Flash-KMeans как раз решает эту задачу: она получает точные результаты k-means с заметно меньшим потреблением памяти и более быстрой сходимостью. Давайте разберемся, что делает k-means трудным в масштабе и как современные варианты с этим справляются.

Как на самом деле работает стандартный k-means

Алгоритм Ллойда, который обычно и подразумевают под «k-means», на каждой итерации делает два шага. Шаг назначения: для каждой точки вычисляем расстояние до всех K центроидов и относим её к ближайшему. Шаг обновления: пересчитываем каждый центроид как среднее всех точек, которые к нему отнесены.

import numpy as np
def kmeans_lloyd(X, K, max_iter=100):
n, d = X.shape
# Random initialization (bad, but we'll fix this later)
centroids = X[np.random.choice(n, K, replace=False)]
for _ in range(max_iter):
# Assignment: O(n * K * d) — this is the bottleneck
distances = np.linalg.norm(X[:, None] - centroids[None, :], axis=2)
labels = np.argmin(distances, axis=1)
# Update: O(n * d)
new_centroids = np.array([
X[labels == k].mean(axis=0) for k in range(K)
])
if np.allclose(centroids, new_centroids):
break
centroids = new_centroids
return labels, centroids

Шаг назначения стоит O(n × K × d) на итерацию, где n — количество точек, K — число кластеров, d — размерность. Для n = 1 миллиард, K = 1000, d = 128 (реалистичный сценарий кластеризации эмбеддингов) это 128 триллионов операций с плавающей запятой за итерацию. Даже при производительности в один терафлопс это 128 секунд на итерацию, а k-means обычно сходится за 10–50 итераций.

Проблема инициализации

Прежде чем k-means сделает хотя бы одну итерацию, нужно выбрать начальные позиции центроидов. Этот выбор важнее, чем думает большинство людей: плохая инициализация может привести к решению, которое произвольно хуже оптимального.

Случайная инициализация (выбор K случайных точек данных в качестве начальных центроидов) проста, но ненадежна. Если два начальных центроида попадут в один кластер, другой кластер останется без представителя. Алгоритм сойдется, но к неоптимальному решению. При K = 100 вероятность хотя бы одной такой коллизии высока.

K-means++ (Arthur и Vassilvitskii, 2007) — стандартное решение. Первый центроид выбирается случайно, а каждый следующий — с вероятностью, пропорциональной квадрату расстояния до ближайшего уже существующего центроида. Точки, далекие от имеющихся центроидов, выбираются с большей вероятностью, и начальные центроиды равномерно распределяются по данным. Это дает гарантию аппроксимации O(log K): решение доказуемо находится в пределах O(log K) от оптимального.

def kmeans_plus_plus_init(X, K):
n, d = X.shape
centroids = [X[np.random.randint(n)]]
for _ in range(1, K):
# Distance from each point to nearest existing centroid
dists = np.min([
np.sum((X - c) ** 2, axis=1) for c in centroids
], axis=0)
# Sample proportional to distance squared
probs = dists / dists.sum()
next_idx = np.random.choice(n, p=probs)
centroids.append(X[next_idx])
return np.array(centroids)

Подвох в том, что сама инициализация k-means++ стоит O(n × K × d) — нужно K проходов по всему датасету. При большом K и большом n инициализация занимает больше времени, чем несколько итераций алгоритма Ллойда. Масштабируемые варианты, например k-means|| (Bahmani и др., 2012), уменьшают число проходов: кандидаты отбираются с избытком параллельно, а затем сводятся к K центроидам.

Стена памяти

Стандартному k-means нужен весь датасет в памяти одновременно. На шаге назначения нужно обойти каждую точку и сравнить её с каждым центроидом. Для миллиарда 128-мерных векторов float32 одни только данные занимают 512 ГБ — больше, чем есть в большинстве отдельных машин.

Наивное решение — mini-batch k-means: на каждой итерации берем случайную подвыборку точек и обновляем центроиды по ней. Это работает, но сходится медленнее и к немного другому (обычно чуть худшему) решению, чем точный k-means. Во многих задачах потеря качества приемлема. Но в других — например, при построении кодовых книг для квантования в векторном поиске — разница в качестве имеет значение.

Flash-KMeans идет другим путем. Вместо того чтобы держать весь датасет в памяти, он использует потоковый проход с умным кешированием. Ключевое наблюдение: большинство точек между итерациями не меняют кластер. Если точка уверенно лежит в кластере 7 (гораздо ближе к центроиду 7, чем к любому другому), проверка всех K расстояний — пустая трата работы. Поддерживая границы расстояний, Flash-KMeans может пропускать полный расчет для большинства точек на большинстве итераций.

Оптимизации на основе неравенства треугольника

Алгоритм Элкана (2003) использует неравенство треугольника, чтобы пропускать ненужные вычисления расстояний. Неравенство говорит: расстояние от точки P до центроида A не превышает суммы расстояния от P до центроида B и расстояния от B до A. Если известно, что P сейчас приписана к центроиду B, и известно расстояние между центроидами A и B, иногда можно доказать, что A слишком далеко от P, не вычисляя реального расстояния.

На практике это исключает 80–95% вычислений расстояний после первых нескольких итераций, когда большинство точек уже близки к своему правильному центроиду. Оставшиеся вычисления касаются точек у границ кластеров — только они могут реально сменить назначение.

Цена вопроса: алгоритму Элкана нужно дополнительно O(n × K) памяти для хранения границ расстояний (нижних границ от каждой точки до каждого центроида и верхних границ до назначенного центроида). При большом K эта память может стать существенной. Алгоритм Хэмерли снижает затраты до O(n), храня для каждой точки только одну нижнюю границу, но и отсекает меньше вычислений.

Ускорение на GPU

Шаг назначения в k-means — задача, которая идеально параллелится: вычисление расстояния для каждой точки не зависит от остальных. Поэтому k-means отлично ложится на GPU. Библиотека cuML от NVIDIA и FAISS от Facebook включают GPU-реализации k-means, которые на больших датасетах дают ускорение в 10–50 раз по сравнению с CPU.

Загвоздка в том, что память GPU ограничена. У A100 80 ГБ — этого хватает примерно на 150 миллионов 128-мерных векторов. Для больших датасетов нужно либо разбивать данные между несколькими GPU, либо использовать потоковый подход: данные загружаются чанками, обрабатываются на GPU, а результаты агрегируются на CPU.

# Using FAISS for GPU-accelerated k-means
import faiss
import numpy as np
# 10 million 128-dimensional vectors
n, d, K = 10_000_000, 128, 1000
X = np.random.randn(n, d).astype('float32')
# CPU k-means (for comparison)
kmeans_cpu = faiss.Kmeans(d, K, niter=20, verbose=True)
kmeans_cpu.train(X)  # ~120 seconds
# GPU k-means (single GPU)
kmeans_gpu = faiss.Kmeans(d, K, niter=20, verbose=True, gpu=True)
kmeans_gpu.train(X)  # ~8 seconds — 15x faster

Когда k-means — неподходящий выбор

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

• Сферические кластеры. K-means предполагает, что кластеры примерно сферические и одинакового размера. Если кластеры вытянутые, неправильной формы или сильно различаются по размеру, k-means будет дробить большие и склеивать маленькие. Гауссовы смеси (GMM) справляются с эллиптическими кластерами, DBSCAN — с произвольными формами.
• Известное K. Число кластеров нужно задать заранее. Если K неизвестно, придется запускать k-means с разными значениями и выбирать лучшее по метрике (силуэт, метод локтя). Это умножает общие вычислительные затраты на количество перебираемых значений K.
• Евклидово расстояние. По умолчанию k-means использует евклидово расстояние. Для текстовых эмбеддингов обычно больше подходит косинусное сходство. Обойти это можно L2-нормализацией векторов (тогда евклидово расстояние эквивалентно косинусному), но об этом легко забыть.
• Чувствительность к выбросам. Одиночный выброс далеко от всех кластеров тянет назначенный ему центроид к себе. Устойчивые варианты, например k-medoids (использует медиану вместо среднего), лучше справляются с выбросами, но обходятся дороже.

Практические советы

Поработав с k-means на датасетах от тысяч до миллиардов точек, я понял, что важнее всего:

1. Всегда используйте инициализацию k-means++. Случайная инициализация не стоит риска. Разница в итоговом качестве кластеров часто составляет 10–30%, а k-means++ добавляет ничтожные накладные расходы для малых и средних K.
2. Запускайте несколько раз. K-means находит локальный оптимум, а не глобальный. Запустите его 5–10 раз с разными случайными seed и выберите лучший результат (с минимальной суммарной внутрикластерной дистанцией). Это дешевая страховка от неудачных запусков.
3. Нормализуйте признаки. Если один признак лежит в диапазоне [0, 1000000], а другой в [0, 1], высокодиапазонный признак будет доминировать в расчете расстояний. Стандартизируйте (вычтите среднее, поделите на стандартное отклонение) или выполните min-max нормализацию перед кластеризацией.
4. Используйте FAISS для крупномасштабной кластеризации. Если точек больше миллиона, k-means из scikit-learn будет медленным. Реализация FAISS сильно оптимизирована и поддерживает GPU из коробки.
5. Для исследовательской работы рассмотрите приближенные методы. Mini-batch k-means работает примерно в 10 раз быстрее точного k-means и обычно дает результаты, достаточно хорошие для исследования. Для продакшен-кластеризации, где важно качество, используйте точный k-means.

K-means — из тех алгоритмов, которыми легко пользоваться, но трудно пользоваться хорошо, и которые стоит понимать глубоко. Разрыв между наивной и оптимизированной реализацией — и по скорости, и по качеству результата — огромен. На малых масштабах ничего из этого не важно. На масштабах, где это важно, понимание инициализации, управления памятью, отсечения расстояний и GPU-ускорения — это разница между кластеризацией, которая занимает часы, и кластеризацией, которая занимает минуты.