Агломеративная кластеризация: алгоритм, быстродействие, код на GitHub

Несколько лет назад мне потребовалось очень качественно кластеризовать относительно неплотные графы среднего размера (сотни тысяч объектов, сотни миллионов связей). Тогда оказалось, что алгоритма с подходящим набором свойств просто не существует, несмотря на всё разнообразие методов, придуманных человечеством за многие десятилетия. Имеющиеся решения работали либо просто очень плохо, либо очень плохо и к тому же медленно.

К счастью, оказалось, что идеи, заложенные в метрики качества кластеризации, о которых я рассказывал в прошлой статье, можно адаптировать и создать на их основе алгоритм кластеризации. Он достигает очень высоких показателей качества и к тому же работает очень быстро благодаря некоторым удачным аналитическим свойствам оптимизируемых величин. Алгоритм относится к классу агломеративных: основной операцией является слияние нескольких уже имеющихся кластеров в один более крупный кластер.

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

1. Агломеративная кластеризация

Пусть даны множество объектов $$ и функция попарной близости объектов $$. Каждый элемент близок сам себе, а функция близости симметрична:

Также определено некоторое разбиение множества $$ на кластеры:

Предполагается, что кластеризация образует покрытие исходного множества. При необходимости можно считать, что изолированные (не принадлежащие ни одному из кластеров) элементы образуют тривиальные кластеры из одного элемента.

Множество кластеров будем обозначать $$:

Пусть выбраны объект $$ и множество $$. Воспользуемся интуицией BCubed-метрик качества кластеризации и определим для этой пары значения точности и полноты:

Заметим, что эти определения удобны тем, что $$ не обязательно находится во множестве $$.

Теперь можно вычислить и средние значения точности и полноты для элементов $$ относительно самого $$:

Определим композицию этих показателей, воспользовавшись идеей метрики ECC:

Здесь параметр $$, как и в определении $$-меры, задаёт относительную важность полноты относительно точности. В простейшем случае $$ и показатели равноправны.

Рассмотрим два непересекающихся подмножества $$ и $$. Если они являются кластерами, то средние точность и полнота входящих в них элементов равны:

Определим и композицию метрик для этого случая:

Если объединить элементы множеств $$ и $$ в один кластер, показатели станут равны:

А композиция метрик определится просто как $$.

При объединении кластеров $$ и $$ композиция метрик для входящих в них элементов изменится на величину

Идея нашего алгоритма агломеративной кластеризации в том, чтобы на каждом шаге объединять два кластера, на которых эта разница достигает максимума. Тогда алгоритм запишется следующим образом:

Даны:

• Множество объектов $$
• Симметричная попарной близости $$
• Тривиальное множество кластеров: $$, $$

Основной цикл:

1. Если $$, вернуть $$ и завершить алгоритм.
2. Выбрать $$.
3. Если $$, вернуть $$ и завершить алгоритм.
4. В противном случае положить $$ и вернуться к шагу 1.

Основной проблемой такого алгоритма будет быстродействие. Каждая итерация основного цикла требует вычисления функции $$ для всех пар кластеров. Но оказывается, что эту часть алгоритма можно существенно ускорить.

2. Ускорение алгоритма

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

Ускорить алгоритм в связи с первым аспектом достаточно просто. Заметим, что после объединения кластеров $$ и $$ значения функции $$ нужно обновить лишь на тех парах кластеров, в которые входит либо $$, либо $$.

Другими словами, величина $$ не меняется после объединения кластеров $$ и $$, если $$. Следовательно, обновится лишь линейное, а не квадратичное по размеру $$ число значений.

Хранить их при этом можно в любой set-подобной структуре, так что извлечение максимального элемента будет занимать $$ времени.

Справиться со вторым аспектом несколько сложнее. Чтобы лучше понять, как здесь работает ускорение, рассмотрим произвольную аддитивную по ячейкам матрицы функцию $$. Пример такой функции — простая сумма значений.

Пусть множество $$ разбито на три кластера. Тогда функция $$ определена на всех девяти прямоугольных фрагментах матрицы, которые соответствуют парам различных кластеров. Нас интересует, как обновить её значения при объединении двух кластеров в один.

После такой операции останется четыре пары кластеров и четыре интересующих нас значения функции. В силу её аддитивности эти четыре значения вычисляются из предшествующих за константное время, не зависящее от размеров кластеров. Правила такого пересчёта показаны на картинке выше.

Теперь заметим: если в ячейках матрицы находятся близости элементов, а функция $$ осуществляет простое суммирование, мы получаем правила для вычисления и обновления характеристик точности. Средние показатели метрики точности для элементов из $$ до и после объединения будут равны, соответственно:

Аналогично можно поступить и с полнотой. Определим аддитивную функцию $$ так, что она будет суммировать нормированные по строкам близости из исходной матрицы. Таким образом:

Сумма значений $$ по элементам одной строки матрицы всегда равняется единице. Тогда:

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

Аналогично можно вычислить показатели, необходимые для определения величины $$:

То есть благодаря аддитивности каждой введённой величины по элементам матрицы все вычисления работают за константное время! Обновление всех значений функции в худшем случае потребует $$ времени.

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

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

4. Реализация

Описанный алгоритм реализован на C++ и выложен под лицензией MIT на GitHub. Реализация компилируется и запускается на Linux и Windows.

Программу легко собрать:

git clone https://github.com/yandex/agglomerative_clustering/
cmake .
cmake --build .

После этого её можно запустить на одном из примеров, находящихся в том же проекте:

./agglomerative_clustering < data/2d_similarities.tsv > 2d_clusters.tsv

В результате будет сгенерирован файл, формат которого совпадает с форматом, в котором описывается разметка кластеризации (см. статью про метрики кластеризации)

5. Наборы данных

Первый поставляемый с программой набор данных — синтетический. 18 343 точки на плоскости получены из 1000 нормальных двумерных распределений со случайными целочисленными центрами. Близость между двумя точками определяется по формуле:

Сами точки из этого набора можно найти в файле data/2d_points.tsv, а можно сгенерировать заново, запустив скрипт data/2d_gen.py. Эталонной считается кластеризация в файле data/2d_markup.tsv, где каждая точка относится к породившему её распределению.

^{_{Синтетический набор точек на плоскости}}

Алгоритм кластеризации должен быть устойчив к потере информации о связях между объектами. Чтобы продемонстрировать это свойство, в наборе из всех близостей, существенно превосходящих ноль, была оставлена примерно одна треть (228 018 штук). Увеличивая параметр -f (который соответствует параметру $$ в этой статье), можно добиться практически тех же показателей качества, что и на графе до удаления рёбер:

./agglomerative_clustering -f 10 < data/2d_similarities.tsv > 2d_clusters.tsv
../cluster_metrics/cluster_metrics data/2d_markup.tsv 2d_clusters.tsv
ECC   0.62411 (0.62411)
BCP   0.74112 (0.74112)
BCR   0.68095 (0.68095)
BCF1  0.70976 (0.70976)

Здесь для вычисления метрик используется программа из предыдущей статьи.

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

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

cmake -DCMAKE_BUILD_TYPE=release .
cmake --build .
tar zxvf data/doc.similarities.tar.gz
./agglomerative_clustering < doc.similarities.tsv > doc.clusters.tsv

Программа в процессе работы выводит информацию о времени выполнения разных этапов:

loaded 5000000 pairs
loading documents: finished in 11.323s
clustering documents: finished in 28.653s
printing clusters: finished in 0.038s

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