k-средних

В фотографии озера Иссык сто пятьдесят тысяч пикселей и десятки тысяч разных цветов. Можно ли оставить шесть так, чтобы картинка осталась той же — с той же тёмной водой, тем же светлым небом? Нужно найти шесть цветов, которые лучше всего представляют все остальные, а для этого — разбить пиксели на шесть групп похожих. Это и делает k-средних. Справа — та же фотография, раскрашенная только найденными цветами; внизу — облако её пикселей в кубе RGB, через которое ползут центры групп.

Нажмите «Новый старт» несколько раз. Иногда центры расходятся по облаку аккуратно, а иногда два из них делят одну и ту же синеву неба, зато все оттенки зелени достаются одному. Алгоритм каждый раз останавливается — но не всегда в лучшем месте. Ниже разберём, почему так выходит и что с этим делать.

Идея

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

1. Отнести. Каждой точке — номер ближайшего центра.
2. Пересчитать. Каждый центр — в среднее арифметическое своих точек.

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

Почему он останавливается

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

Различных разбиений n точек на k групп конечное число, и, пока разбиение меняется, инерция строго уменьшается — к одному разбиению алгоритм не возвращается. Поэтому рано или поздно разбиение перестанет меняться, и алгоритм остановится. На практике это десятки шагов, на фото озера — около двадцати.

Но остановка — это не победа. Алгоритм находит локальный минимум: такое разбиение, которое нельзя улучшить ни одним из двух шагов. Лучшего разбиения он не ищет, и найти его в общем случае трудно: задача k-средних NP-трудна.

Всё решает старт

Вот двенадцать точек в трёх ясных группах по углам. Начнём нарочно плохо: все три центра — в одном углу.

step 1: inertia 171.05, centres (1.00, 1.00) (4.39, 7.24) (2.00, 1.20)
step 2: inertia 93.94, centres (1.00, 1.00) (4.99, 8.50) (2.00, 1.87)
labels [0, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1]

Алгоритм честно уменьшает инерцию и останавливается на 93,94: две верхние группы слиплись в одну, а нижняя разделена на точку и три. Ни один шаг уже ничего не меняет. А если начать с трёх точек из разных углов, хватает одного шага, и инерция — 9,85, почти в десять раз меньше.

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

Код

Обе части шага прямо как в определении. При равенстве расстояний точка уходит к центру с меньшим номером, а центр пустой группы остаётся на месте — такие договорённости делают результат воспроизводимым.

А это палитра настоящей фотографии — той самой, с озером. Программа читает картинку в формате PPM (простой формат: заголовок и байты R, G, B подряд), раскладывает её 9600 пикселей на шесть цветов и печатает палитру с долями. Python и Go дают одни и те же цвета до последней цифры: старт один и тот же, и арифметика с плавающей точкой у них общая.

Сколько это стоит

Шаг «отнести» сравнивает каждую из n точек с каждым из k центров, и сравнение стоит O(d) для точек размерности d: всего O(n·k·d). Шаг «пересчитать» — O(n·d). Шагов обычно немного, от единиц до нескольких десятков, хотя в худшем случае их может понадобиться экспоненциально много — но такие примеры приходится строить специально.

Для фотографии в 150 000 пикселей и 6 цветов шаг — около миллиона операций, мгновение. Для миллионов точек в сотнях измерений уже не мгновение, и тогда берут мини-батчи: каждый шаг считают по случайной части точек. В сцене так и сделано — центры подбираются по 12 000 случайным пикселям, а раскрашивается вся картинка.

Сколько групп брать

Алгоритм не знает, сколько групп «на самом деле». Инерция всегда падает с ростом k: при k = n она равна нулю, каждая точка — сама себе группа. Поэтому смотрят, где падение замедляется — так называемый «локоть» графика. Для палитры вопрос проще: k — это сколько цветов вы готовы хранить. Подвигайте ползунок k: на двух цветах озеро становится графикой, на шести узнаётся, на двенадцати почти не отличается от оригинала.

Где k-средних ошибается

• Группы не круглые. Алгоритм режет пространство на области ближайшего центра — выпуклые куски. Вытянутые, изогнутые или вложенные друг в друга группы он разрежет неправильно. Для них есть DBSCAN, который ищет плотные области любой формы.
• Масштаб признаков. Если одна координата — рост в сантиметрах, а другая — возраст в годах, расстояние почти целиком определит рост. Признаки сначала приводят к одному масштабу.
• Цвет — не RGB. Евклидово расстояние в RGB не совпадает с тем, насколько цвета различны для глаза. Для палитр лучше работать в пространстве Lab, где расстояния ближе к восприятию.
• Выбросы. Одна далёкая точка тянет среднее за собой. Если выбросов много, берут медианы вместо средних (k-медиан) или настоящие точки вместо средних (k-медоидов).

История

Идею разбиения по ближайшим центрам предложил Гуго Штейнгауз в 1956 году. Алгоритм с двумя чередующимися шагами придумал Стюарт Ллойд в Bell Labs в 1957 году — для квантования сигнала в импульсно-кодовой модуляции, то есть ровно для задачи «какие уровни оставить, чтобы сигнал исказился меньше всего». Опубликовал он его только в 1982 году. Название «k-means» появилось в работе Джеймса Маккуина 1967 года.

Задачи

Первая задача — половина шага: отнести пиксели к ближайшим цветам палитры. Вторая — весь алгоритм с точными правилами, которые делают ответ единственным.