Агломеративная кластеризация: алгоритм, быстродействие, код на GitHub11.06.2020 09:32

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

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

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

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

Пусть даны множество объектов $A=\{a_1, ..., a_n\}$ и функция попарной близости объектов $s:A^2 \rightarrow [0,1]$ . Каждый элемент близок сам себе, а функция близости симметрична:

$s(a_i,a_i) = 1$

$s(a_i, a_j) = s(a_j, a_i)$

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

$A=C_1 \sqcup C_2 \sqcup ... \sqcup C_k$

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

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

$\mathcal{L} = \{C_1, C_2, ..., C_k\}$

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

${P}(x, X) = \frac{1}{|X|}\sum_{x' \in X}{s(x, x')}$

${R}(x) = \frac{\sum_{x' \in X}{s(x, x')}}{\sum_{x' \in A}{s(x, x')}}$

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

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

$P(X) = \frac{1}{|X|}\sum_{x \in X}P(x, X)$

$R(X) = \frac{1}{|X|}\sum_{x \in X}R(x, X)$

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

$M_\alpha(X) = R^\alpha(X) \cdot P(X)$

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

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

$P(X, Y) = \frac{1}{|X \cup Y|}\Big[ \sum_{x \in X}{P(x, X)} + \sum_{y \in Y}{P(y, Y)}\Big]$

$R(X, Y) = \frac{1}{|X \cup Y|}\Big[ \sum_{x \in X}{R(x, X)} + \sum_{y \in Y}{R(y, Y)}\Big]$

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

$M_\alpha(X, Y) = R^\alpha(X, Y) \cdot P(X, Y)$

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

$P(X \cup Y) = \frac{1}{|X \cup Y|}\sum_{a \in X \cup Y}{P(a, X \cup Y)}$

$R(X \cup Y) = \frac{1}{|X \cup Y|}\sum_{a \in X \cup Y}{R(a, X \cup Y)}$

А композиция метрик определится просто как $M_\alpha(X \cup Y)$ .

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

$\Delta_{\alpha}{X, Y} = M_\alpha(X \cup Y) - M_\alpha(X, Y)$

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

Даны:
Множество объектов $A = \{a_1, a_2, ..., a_n \}$

Симметричная попарной близости $s:A^2 \rightarrow [0,1]$

Тривиальное множество кластеров: $\mathcal{L} = \{C_1, C_2, ..., C_n\}$ , $C_i = \{a_i\}$

Основной цикл:
Если $|\mathcal{L}| = 1$ , вернуть $\mathcal{L}$ и завершить алгоритм.

Выбрать $C_i, C_j = \arg \max_{X, Y \in \mathcal{L}} \Delta_{\alpha}(X, Y)$ .

Если $\Delta_{\alpha}(C_i, C_j) < 0$ , вернуть $\mathcal{L}$ и завершить алгоритм.

В противном случае положить $\mathcal{L} = \mathcal{L} \setminus \{C_i, C_j\} \cup \{C_i \cup C_j\}$ и вернуться к шагу 1.

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

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

Сложность алгоритма проявляется в двух аспектах. Первый связан с тем, что количество пар кластеров на шаге 2 является квадратичным от общего числа имеющихся кластеров; второй — с тем, что вычисление величины $\Delta_{\alpha}(C_i, C_j)$ при наивной реализации требует времени, пропорционального как минимум $|C_i| \cdot |C_j|$ .

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

Другими словами, величина $\Delta_{\alpha}(C_k, C_l)$ не меняется после объединения кластеров $C_i$ и $C_j$ , если $\{i, j\} \cap \{k, l\} = \varnothing$ . Следовательно, обновится лишь линейное, а не квадратичное по размеру $\mathcal{L}$ число значений.

Хранить их при этом можно в любой set-подобной структуре, так что извлечение максимального элемента будет занимать $O(\ln{|\mathcal{L}|}^2)=O(\ln|\mathcal{L}|)$ времени.

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

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

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

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

$P(X,Y) = \frac{1}{|X \cup Y|}\Big[\frac{f(X,X)}{|X|} + \frac{f(Y,Y)}{|Y|}\Big]$

$P(X \cup Y) = \frac{f(X,X) + f(X,Y) + f(Y,X) + f(Y,Y)}{|X \cup Y|^2}$

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

$g(x,y) = \frac{s(x,y)}{\sum_{a \in A}s(x, a)}$

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

$R(X,Y) = \frac{g(X,X) + g(Y,Y)}{|X \cup Y|}$

$R(X \cup Y) = \frac{g(X,X) + g(X,Y) + g(Y,X) + g(Y,Y)}{|X \cup Y|}$

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

Аналогично можно вычислить показатели, необходимые для определения величины $\Delta_\alpha(X \cup Y, Z)$ :

$P(X \cup Y,Z) = \frac{1}{|X \cup Y \cup Z|}\Big[\frac{f(X \cup Y,X \cup Y)}{|X \cup Y|} + \frac{f(Z,Z)}{|Z|}\Big]$

$P(X \cup Y \cup Z) = \frac{f(X \cup Y,X \cup Y) + f(X \cup Y,Z) + f(Z,X \cup Y) + f(Z,Z)}{|X \cup Y \cup Z|^2}$

$R(X \cup Y,Z) = \frac{g(X \cup Y,X \cup Y) + g(Z,Z)}{|X \cup Y \cup Z|}$

$R(X \cup Y \cup Z) = \frac{g(X \cup Y,X \cup Y) + g(X \cup Y,Z) + g(Z,X \cup Y) + g(Z,Z)}{|X \cup Y \cup Z|}$

То есть благодаря аддитивности каждой введённой величины по элементам матрицы все вычисления работают за константное время! Обновление всех значений функции в худшем случае потребует $O(|\mathcal{L}| \cdot \ln|\mathcal{L}|)$ времени.

Поскольку каждая итерация уменьшает число кластеров на единицу, полное выполнение алгоритма потребует $O(n^2\ln n)$ времени. Если в среднем по всем итерациям один кластер оказывается связан с $k$ другими, оценка временной сложности улучшается до $O(n \cdot k \cdot \ln n)$ .

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

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 нормальных двумерных распределений со случайными целочисленными центрами. Близость между двумя точками определяется по формуле:

$s(p_1, p_2) = \frac{2}{1 + \exp\Big(\|p_1-p_2\|_2^2\Big)}$

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

Синтетический набор точек на плоскости

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

./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 элементов с пятью миллионами связей занимает около полуминуты на моей виртуальной машине. Такого быстродействия более чем достаточно в приложениях, для которых разрабатывался алгоритм. Программу можно дополнительно ускорить за счёт параллелизма и хардкорных оптимизаций, но рассказ о них тянет на отдельный пост.