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


Задача


Множество точек $X= \{ x_i, i\in1..N \}$ нужно разбить на $K$ кластеров.


Идея решения


Составим вероятностную модель распределения точек по кластерам. Найдём параметры модели, при которых вероятность наблюдать множество $X$ максимальна. С этими параметрами мы сможем определять, к какому кластеру наиболее вероятно относится данная точка $x$.


Модель данных


Введем ряд обозначений, заимствованных из курса.


$p(x)$ — вероятность наблюдать точку $x$.


$p(X) = \prod_{i=1}^{N}p(x_i)$ — вероятность наблюдать множество $X$.


$p_j (x) = \varphi(x; \theta_j)$ — вероятность встретить точку $x$ в кластере $j$. Это распределение параметризовано параметром (или вектором параметров) $\theta_j$, индивидуальным для кластера $j$.


$w_j$ — вероятность кластера $j$, т.е. вероятность того, что случайно выбранная точка относится к кластеру $j$. Случайно выбранная точка точно относится к какому-то кластеру, поэтому $\sum_{j=1}^K w_j = 1$.


Из определений выше следует, что $p(x) = \sum_{j=1}^K w_j p_j(x) = \sum_{j=1}^K w_j \varphi(x; \theta_j)$, т.е. распределение точек моделируется как смесь распределений кластеров.


В итоге, вероятностная модель множества точек $X$:


$p(X) = \prod_{i=1}^{N}\left(\sum_{j=1}^K w_j \varphi(x_i; \theta_j)\right)$


Поиск параметров


Параметры модели $w$ и $\theta$, как и обсуждалось выше, должны обеспечить максимальную вероятность нашим данным:


$w, \theta = \textrm{argmax} \ p(X) = \textrm{argmax} \ \log p(X) = \textrm{argmax}_{w, \theta} \sum_{i=1}^{N} \log \left(\sum_{j=1}^K w_j \varphi(x_i; \theta_j)\right)$


Сумма под знаком логарифма мешает решить задачу аналитически. Ограничение $\sum_{j=1}^K w_j = 1$ не дает нам просто применить автоматическое дифференцирование (например, TensorFlow или PyTorch).


Придется использовать трюк:


L := нижняя граница log p(X)
while log p(X) увеличивается заметно:
    приближаем L к log p(X)
    w, theta = argmax L

Иначе говоря, мы не будем оптимизировать сам $\log p(X)$, раз это сложно. Мы возьмем его нижнюю границу $\mathcal{L}$ и будем итеративно делать два шага:


  1. Уточнять $\mathcal{L}$: "подтягивать" её вверх, ближе к $\log p(X)$.
  2. Искать $w$ и $\theta$, максимизирующие $\mathcal{L}$.

Мы надеемся, что если полученные параметры максимизируют "близкую" нижнюю границу функции, то они неплохо максимизируют и саму функцию.


Нижняя граница $\mathcal{L}$


Ограничим снизу выражение:


$ \log p(X) = \sum_{i=1}^{N} \log \left(\sum_{j=1}^K w_j \varphi(x_i; \theta_j)\right)$


Сначала дополним нашу вероятностную модель. Введем распределение вероятностей $g_i$ данной точки $x_i$ быть в том или ином кластере:


$g_i(j) \equiv p(\textrm{быть в кластере} \ j| \textrm{это точка} \ i)$


Домножим и поделим на это распределение слагаемые под логарифмом:


$\sum_{i=1}^{N} \log \left(\sum_{j=1}^K w_j \varphi(x_i; \theta_j)\right) =\sum_{i=1}^{N} \log \left(\sum_{j=1}^K \frac{ g_i(j) }{ g_i(j) } w_j \varphi(x_i; \theta_j)\right)$


Теперь применим неравенство Йенсена. Оно позволяет логарифм взвешенной суммы ограничить снизу взвешенной суммой логарифмов:


$\log \left(\sum_i q_i x_i \right) \geq \sum_i q_i \log x_i$


Чтобы неравентсво выполнялось, нужно чтобы веса $q_i$ были неотрицательны и их сумма была $1$.


В нашем случае $g_i(j)$ будет весом: его значения неотрицательны и $\sum_{j=1}^K g_i(j) = 1$. Применим неравенство:


$\sum_{i=1}^{N} \log \left(\sum_{j=1}^K \frac{ g_i(j) }{ g_i(j) } w_j \varphi(x_i; \theta_j)\right) \geq \sum_{i=1}^{N} \sum_{j=1}^K g_i(j) \log \left(\frac{ w_j \varphi(x_i; \theta_j) }{ g_i(j) }\right)$


Полученное выражение и будем использовать в качестве нижней границы:


$\mathcal{L}(g, w, \theta) \equiv \sum_{i=1}^{N} \sum_{j=1}^K g_i(j) \log \left(\frac{ w_j \varphi(x_i; \theta_j) }{ g_i(j) }\right)$


Уточняем $\mathcal{L}$ (E-шаг)


Попробуем максимально приблизить $\mathcal{L}(g, w, \theta)$ к $\log p(X)$. Параметры $w$ и $\theta$ будем считать фиксированными, а приближать $\mathcal{L}$ будем через выбор распределения $g$.


Запишем разность $\log p(X)$ и $\mathcal{L}$, а потом посмотрим, как её уменьшить:


$\log p(X) - \mathcal{L}(g, w, \theta) = \sum_{i=1}^N \log p(x_i) - \sum_{i=1}^{N} \sum_{j=1}^K g_i(j) \log \left(\frac{ w_j \varphi(x_i; \theta_j) }{ g_i(j) }\right)=$


$= \sum_{i=1}^N \left(\log p(x_i) \sum_{j=1}^K g_i(j) - \sum_{j=1}^K g_i(j) \log \frac{w_j \varphi(x_i; \theta_j)}{g_i(j)} \right) = \sum_{i=1}^N \sum_{j=1}^K g_i(j) \log \frac{p(x_i) g_i(j)}{w_j \varphi(x_i; \theta_j)}$


Заметим, что под знаком логарифма можно выделить апостериорную вероятность кластера $j$:


$p(j|x_i) = \frac{\varphi(x_i; \theta_j) w_j}{p(x_i)}$


Тогда:


$\sum_{i=1}^N \sum_{j=1}^K g_i(j) \log \frac{p(x_i) g_i(j)}{w_j \varphi(x_i; \theta_j)} = \sum_{i=1}^N \sum_{j=1}^K g_i(j) \log \frac{g_i(j)}{p(j|x_i)}= \sum_{i=1}^N \mathbb{E}_{g_i} \frac{g_{i}}{p(j|x_i)}$


Мы получили интересное выражение: матожидание отношения двух распределений по первому из них. Оно называется дивергеницией Кульбака-Лейблера (или кратко KL-дивергенцией) и служит "расстоянием" между вероятностными распределениями.


Таким образом, разность $\log p(X)$ и $\mathcal{L}$ — это сумма KL-дивергенций:


$\log p(X) - \mathcal{L}(g, w, \theta) = \sum_{i=1}^N KL(g_i || p(j|x_i)) $


KL-дивергенция неотрицательна, поэтому лучшее, что мы можем сделать для приближения нижней границы — это сделать KL-дивергенцию нулевой. А этого несложно добиться: KL-дивергенция будет нулём, если её аргументы — это одинаковые распределения. Вот и сделаем распределение $g_i(j)$ тождественным $p(j|x_i)$:


$g_i(j) = p(j|x_i) = \frac{w_j \varphi(x_i; \theta_j)}{p(x_i)}$


Вычисление $g_i(j)$ по данной формуле и позволит нам приблизить нижнюю границу $\mathcal{L}$ к $\log p(X)$.


Максимизируем $\mathcal{L}$ по параметрам (M-шаг)


Теперь вторая часть итерации: поиск параметров по нижней границе. В этой части наши предположения будут противоположными:


  • распределение $g$ фиксированно;
  • параметры $w$ и $\theta$ подлежат оптимизации.

Перед оптимизацией немного упростим $\mathcal{L}$:


$\mathcal{L}(g, \theta) = \sum_{i=1}^N\left( \sum_{j=1}^K g_i(j) \log \frac{w_j p(x_i; \theta_j)}{g_i(j)} \right) =$


$ = \sum_{i=1}^N \sum_{j=1}^K g_i(j) \log \left( w_j p(x_i; \theta_j) \right) -\sum_{i=1}^N \sum_{j=1}^K g_i(j) \log g_i(j)$


Второе слагаемое не зависит от параметров $w$ и $\theta$, поэтому далее будем оптимизировать только первое слагаемое:


$w, \theta = \textrm{argmax}_{w, \theta}\sum_{i=1}^N \sum_{j=1}^K g_i(j) \log \left( w_j \varphi(x_i; \theta_j) \right)$


Разложим логарифм произведения на сумму логарифмов и получим:


$w = \textrm{argmax}_{w}\sum_{i=1}^N \sum_{j=1}^K g_i(j) \log w_j, \textrm{ при условии }\sum_{j=1} w_j = 1$


$\theta_j = \textrm{argmax} \sum_{i=1}^N g_i(j) \log \varphi (x_i; \theta_j)$


Первая задача решается методом множителей Лагранжа. Результат:


$w_j = \frac{1}{N} \sum_{i=1}^N g_i(j)$


Решение второй задачи зависит от конкретного вида распределения кластера $\varphi (x_i; \theta_j)$. Как видно, для её решение больше не придётся иметь дело с суммой под знаком логарифма, поэтому, например, для гауссовых распределений решение может быть выписано аналитически.


Итог


Мы выяснили, в чем заключается суть итераций EM-алгоритма для кластеризации, и увидели, как в общем виде выводятся их формулы.