Смеси распределений и EM-алгоритм

Заметка о кластеризации закончилась методами, которые дают жёсткие ответы. Метод k-средних помещает каждую точку ровно в один кластер, ничего не говорит о своей уверенности и ведёт себя так, будто все кластеры круглые и одинаково широкие. Эта заметка превращает кластеризацию в вероятностную модель. Каждый кластер становится гауссианой со своим средним, ковариационной матрицей и весом, данные — выборка из их смеси, а параметры подбираются методом максимального правдоподобия. У правдоподобия нет максимума в замкнутом виде, потому что в него входят принадлежности к кластерам как ненаблюдаемые переменные. EM-алгоритм (EM algorithm) работает с такими моделями, поднимаясь по нижней оценке логарифма правдоподобия, по одному блоку переменных за раз. Та же нижняя оценка, в которой точное апостериорное распределение заменено нейросетью, — целевая функция вариационного автоэнкодера.

Обозначения. $x_1, \dots, x_n \in \mathbb R^d$ — данные, $K$ — число компонент. $z_i \in \{1, \dots, K\}$ — ненаблюдаемая компонента точки $i$. У компоненты есть вес $\pi_k \ge 0$, причём $\sum_k \pi_k = 1$, среднее $\mu_k$ и ковариационная матрица $\Sigma_k$, симметричная положительно определённая; $\theta$ объединяет их все. $\mathcal N(x \mid \mu, \Sigma)$ — гауссовская плотность со средним $\mu$ и ковариацией $\Sigma$ (вторым аргументом идёт дисперсия, а не стандартное отклонение). $\ell(\theta) = \sum_i \log p(x_i \mid \theta)$ — логарифм правдоподобия, $r_{ik}$ — ответственности, определённые ниже, $q$ — распределение на латентных переменных.

Модель гауссовской смеси#

Гауссовская смесь (Gaussian mixture) описывает данные двухшаговой порождающей историей. Сначала выбираем компоненту $z$ с вероятностями $\pi_1, \dots, \pi_K$, затем сэмплируем $x$ из гауссианы этой компоненты:

$$z \sim \text{Categorical}(\pi_1, \dots, \pi_K), \qquad x \mid z = k \;\sim\; \mathcal N(\mu_k, \Sigma_k) .$$

Компонента никогда не наблюдается, поэтому распределение одного $x$ — маргинальное; оно получается суммированием совместного распределения $p(x, z = k) = \pi_k \mathcal N(x \mid \mu_k, \Sigma_k)$ по $k$ по формуле полной вероятности:

$$p(x \mid \theta) = \sum_{k=1}^K \pi_k\, \mathcal N(x \mid \mu_k, \Sigma_k) .$$

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

Разобранный пример. В одном измерении возьмём $\pi = (0.3, 0.7)$, $\mu = (0, 4)$ и стандартные отклонения $1$ и $2$. В точке $x = 2$ плотности двух компонент равны $\mathcal N(2 \mid 0, 1) = 0.0540$ и $\mathcal N(2 \mid 4, 4) = 0.1210$, поэтому

$$p(2) = 0.3 \cdot 0.0540 + 0.7 \cdot 0.1210 = 0.0162 + 0.0847 = 0.1009 .$$

Правдоподобие с латентными переменными#

Для независимых одинаково распределённых данных логарифм правдоподобия равен

$$\ell(\theta) = \sum_{i=1}^n \log \sum_{k=1}^K \pi_k\, \mathcal N(x_i \mid \mu_k, \Sigma_k) .$$

Логарифм суммы не распадается на слагаемые. Приравнивание градиента к нулю связывает каждый параметр со всеми остальными через сумму в знаменателе, и решения в замкнутом виде нет. Сравним с логарифмом правдоподобия полных данных (complete-data log-likelihood) — тем, который был бы у нас, если бы компоненты $z_i$ наблюдались:

$$\log p(x, z \mid \theta) = \sum_{i=1}^n \Big(\log \pi_{z_i} + \log \mathcal N(x_i \mid \mu_{z_i}, \Sigma_{z_i})\Big) .$$

Это сумма логарифмов. Она распадается на подгонку одной гауссианы для каждой компоненты, по точкам с меткой этой компоненты, и одну подгонку категориального распределения для весов; у каждой есть замкнутый вид из заметки о методе максимального правдоподобия: доля точек, среднее и ковариация с множителем $1/n_k$. Вся трудность — в том, что $z$ неизвестны. EM-алгоритм превращает это наблюдение в метод: заменить неизвестные метки их вероятностями при данных, подогнать параметры так, будто метки известны, и повторить.

Три свойства $\ell$ полезно знать до её максимизации. Она не вогнута и имеет несколько локальных максимумов. Она симметрична относительно перенумерации компонент, поэтому каждый максимум встречается в $K!$ копиях. И она не ограничена: компонента с центром в одной точке данных, дисперсия которой стремится к нулю, даёт $\ell \to \infty$ — этот случай разобран в заметке о методе максимального правдоподобия. Поэтому оценки максимального правдоподобия для гауссовской смеси без ограничений не существует: супремум бесконечен, и ни одно значение параметров его не достигает. EM возвращает локальный максимум правдоподобия вдали от вырождающихся решений или, с нижним порогом для ковариаций или с априорным распределением, как в разделе о вырожденных решениях ниже, точку максимума ограниченного или регуляризованного правдоподобия.

Ответственности#

При заданных параметрах вероятность того, что точка $x_i$ порождена компонентой $k$, следует из формулы Байеса: априорная вероятность $\pi_k$, правдоподобие $\mathcal N(x_i \mid \mu_k, \Sigma_k)$, обоснованность $p(x_i \mid \theta)$:

$$r_{ik} = p(z_i = k \mid x_i, \theta) = \frac{\pi_k\, \mathcal N(x_i \mid \mu_k, \Sigma_k)}{\sum_{j=1}^K \pi_j\, \mathcal N(x_i \mid \mu_j, \Sigma_j)} .$$

Числа $r_{ik}$ — это ответственности (responsibilities): каждая строка — распределение по компонентам, $r_{ik} \ge 0$ и $\sum_k r_{ik} = 1$. Это мягкий вариант назначения в k-средних. В разобранном примере выше точка $x = 2$ принадлежит первой компоненте с вероятностью $0.0162/0.1009 = 0.16$, а второй — с вероятностью $0.84$: она лежит в двух стандартных отклонениях от первого среднего и в одном от второго, а вторая компонента к тому же встречается чаще. Метод k-средних, сравнивающий только расстояния, поместил бы $x = 2$ посередине, на границу.

Откуда берётся нижняя оценка#

Возьмём одну точку $x$ с латентной переменной $z$ (для всего набора данных просуммируем по точкам) и пусть $q(z)$ — любое распределение на $K$ значениях $z$, положительное везде, где положительно апостериорное распределение. Поскольку $p(x, z \mid \theta) = p(z \mid x, \theta)\, p(x \mid \theta)$,

$$\log p(x \mid \theta) = \log \frac{p(x, z \mid \theta)}{q(z)} - \log \frac{p(z \mid x, \theta)}{q(z)}$$

для любого $z$. Левая часть не зависит от $z$, поэтому усреднение обеих частей по $z \sim q$ её не меняет:

$$\log p(x \mid \theta) = \underbrace{\sum_z q(z) \log \frac{p(x, z \mid \theta)}{q(z)}}_{\mathcal L(q, \theta)} + \underbrace{\sum_z q(z) \log \frac{q(z)}{p(z \mid x, \theta)}}_{D_{\mathrm{KL}}(q \,\|\, p(z \mid x, \theta))} .$$

Второе слагаемое — расхождение Кульбака–Лейблера (KL divergence) распределения $q$ относительно апостериорного. Оно неотрицательно и равно нулю ровно тогда, когда $q$ совпадает с апостериорным распределением. Следовательно,

$$\log p(x \mid \theta) \;\ge\; \mathcal L(q, \theta) = \mathbb E_{q}\big[\log p(x, z \mid \theta)\big] + H(q), \qquad H(q) = -\sum_z q(z) \log q(z),$$

причём равенство достигается, когда $q$ — апостериорное распределение $p(z \mid x, \theta)$. $\mathcal L$ — это нижняя оценка обоснованности (evidence lower bound, ELBO); Neal and Hinton (1998) называют её со знаком минус свободной энергией. То же неравенство в одну строку следует из неравенства Йенсена для вогнутого логарифма, $\log \mathbb E_q[p(x, z)/q(z)] \ge \mathbb E_q[\log p(x, z)/q(z)]$, но тождество выше говорит больше: зазор в точности равен KL-дивергенции, так что мы знаем, когда оценка точна.

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

EM — покоординатный подъём по нижней оценке#

$\mathcal L(q, \theta)$ зависит от двух блоков: распределения $q$ (для независимых одинаково распределённых данных — своё $q_i$ для каждой точки) и параметров $\theta$. EM-алгоритм (Dempster, Laird and Rubin, 1977) максимизирует её по одному блоку за раз:

1. E-шаг. При фиксированном $\theta^t$ максимизируем по $q$. Левая часть $\log p(x \mid \theta^t)$ не зависит от $q$, поэтому максимизировать $\mathcal L$ — значит минимизировать слагаемое KL, и максимум достигается на апостериорном распределении: $q^{t+1}(z) = p(z \mid x, \theta^t)$. Для смеси это вычисление ответственностей $r_{ik}$. После E-шага нижняя оценка касается логарифма правдоподобия в $\theta^t$.
2. M-шаг. При фиксированном $q^{t+1}$ максимизируем по $\theta$. Энтропия $H(q^{t+1})$ постоянна, поэтому

$$\theta^{t+1} = \arg\max_\theta\, \sum_{i=1}^n \mathbb E_{q_i^{t+1}}\big[\log p(x_i, z_i \mid \theta)\big] ,$$

— ожидаемый логарифм правдоподобия полных данных, в исходной статье он обозначен $Q(\theta, \theta^t)$.

Монотонность следует в одну строку: сначала нижняя оценка, затем M-шаг, затем точность оценки после E-шага:

$$\ell(\theta^{t+1}) \;\ge\; \mathcal L(q^{t+1}, \theta^{t+1}) \;\ge\; \mathcal L(q^{t+1}, \theta^{t}) \;=\; \ell(\theta^{t}) .$$

Логарифм правдоподобия никогда не убывает. Любое $\theta$, увеличивающее нижнюю оценку, а не только точка её максимума, сохраняет неравенство, поэтому M-шаг можно заменить несколькими шагами градиентного подъёма (обобщённый EM, generalized EM). Ещё два факта завершают картину. В $\theta^t$ нижняя оценка $\mathcal L(q^{t+1}, \cdot)$ и $\ell$ касаются и имеют одинаковый градиент, потому что их разность, KL-дивергенция, достигает там минимума 0; поэтому EM перестаёт двигаться только в стационарных точках $\ell$. Неубывание — ещё не сходимость. $\ell(\theta^t)$ сходится к конечному значению, если $\ell$ ограничена сверху вдоль итераций, а для сходимости параметров нужно больше. Wu (1983) показывает, что если множество $\{\theta : \ell(\theta) \ge \ell(\theta^0)\}$ компактно, а $\ell$ и нижняя оценка непрерывны, то каждая предельная точка итераций EM — стационарная точка $\ell$, а $\ell(\theta^t)$ монотонно сходится; стационарная точка может быть локальным максимумом или седлом, и нет гарантии, что это глобальный максимум. Для гауссовской смеси компактность нарушается у вырождающихся решений выше, поэтому гарантия относится только к запускам, которые держатся от них в стороне, например с нижним порогом для ковариаций.

Рисунок

EM как подъём по нижней оценке

Разобранный пример. В модели с рисунка один неизвестный параметр — среднее $\theta$ второй из двух компонент с равными весами и единичной дисперсией; первая зафиксирована, $\mathcal N(0, 1)$. При $\theta_0 = 0.6$ ответственности второй компоненты для семи точек $-0.8, -0.2, 0.3, 0.9, 2.1, 2.8, 3.4$ равны $0.341, 0.426, 0.500, 0.589, 0.746, 0.818, 0.865$: точки далеко справа приписываются в основном ей. При фиксированных весах и первой компоненте M-шаг для $\theta$ — взвешенное среднее:

$$\theta_1 = \frac{\sum_i r_i x_i}{\sum_i r_i} = \frac{7.121}{4.285} = 1.662 .$$

Логарифм правдоподобия растёт с $-16.35$ до $-12.83$, затем до $-12.14$ при $\theta_2 = 2.168$, и EM сходится к максимуму в $\theta = 2.54$ с $\ell = -12.00$. Шаги уменьшаются, потому что вблизи максимума нижняя оценка и логарифм правдоподобия становятся похожи: EM сходится линейно, сначала быстро, а в конце медленно.

Формулы обновления для гауссовских смесей#

Для гауссовской смеси ожидаемый логарифм правдоподобия полных данных равен

$$Q(\theta) = \sum_{i=1}^n \sum_{k=1}^K r_{ik} \Big( \log \pi_k + \log \mathcal N(x_i \mid \mu_k, \Sigma_k) \Big),$$

то есть логарифму правдоподобия полных данных, в котором каждый индикатор $[z_i = k]$ заменён его апостериорной вероятностью $r_{ik}$. Он распадается по параметрам. Обозначим $N_k = \sum_i r_{ik}$ эффективное число точек компоненты $k$.

Веса. Максимизируем $\sum_k N_k \log \pi_k$ при условии $\sum_k \pi_k = 1$. У лагранжиана $\sum_k N_k \log \pi_k + \lambda(\sum_k \pi_k - 1)$ производная $N_k/\pi_k + \lambda = 0$, поэтому $\pi_k \propto N_k$, и после нормировки $\pi_k = N_k / n$: доля точек с дробными счётчиками.

Средние. Слагаемые с $\mu_k$ равны $-\frac12 \sum_i r_{ik} (x_i - \mu_k)^\top \Sigma_k^{-1} (x_i - \mu_k)$. Градиент $\Sigma_k^{-1} \sum_i r_{ik} (x_i - \mu_k)$ обращается в ноль, когда $\sum_i r_{ik}(x_i - \mu_k) = 0$, потому что $\Sigma_k^{-1}$ обратима.

Ковариации. С обратной ковариационной матрицей (precision) $\Lambda_k = \Sigma_k^{-1}$ слагаемые с $\Sigma_k$ равны $\frac{N_k}{2} \log \det \Lambda_k - \frac12 \operatorname{tr}(\Lambda_k S_k)$, где $S_k = \sum_i r_{ik} (x_i - \mu_k)(x_i - \mu_k)^\top$. Так как $\partial \log\det \Lambda / \partial \Lambda = \Lambda^{-1}$ и $\partial \operatorname{tr}(\Lambda S)/\partial \Lambda = S$, градиент $\frac{N_k}{2}\Lambda_k^{-1} - \frac12 S_k$ обращается в ноль при $\Sigma_k = S_k / N_k$.

Вместе M-шаг таков:

$$\pi_k = \frac{N_k}{n}, \qquad \mu_k = \frac{1}{N_k} \sum_{i=1}^n r_{ik}\, x_i, \qquad \Sigma_k = \frac{1}{N_k} \sum_{i=1}^n r_{ik}\, (x_i - \mu_k)(x_i - \mu_k)^\top ,$$

— оценки максимального правдоподобия для гауссианы, в которых каждая точка учитывается с весом $r_{ik}$ (в $\Sigma_k$ используется новое $\mu_k$). Одна итерация EM стоит $O(nKd^2)$ на ответственности плюс $O(Kd^3)$ на обращения матриц и определители. На практике E-шаг вычисляют в логарифмах, устойчиво вычисляя $\log \sum_k \exp(\cdot)$, потому что в высокой размерности плотности обращаются в 0 при потере точности (underflow).

Рисунок

EM для смеси трёх гауссиан

Рисунок нарочно начинает EM плохо: все три средних внутри одного облака. После одной итерации компоненты уже разошлись, после пяти найдены вытянутое и маленькое облака, а к итерации 25 третья компонента обосновалась на широком облаке. Средний логарифм правдоподобия растёт с $-7.020$ до $-3.223$ и ни разу не падает. Эллипсы показывают то, что не мог представить метод k-средних: один наклонный, один маленький и один широкий кластер, каждый со своей ковариационной матрицей.

k-средних — жёсткий предел с равными дисперсиями#

Зафиксируем все ковариационные матрицы, $\Sigma_k = s^2 I$ с известным $s$, и будем обновлять только средние (веса можно фиксировать или обновлять, предел от них не зависит). Ответственности становятся такими:

$$r_{ik} = \frac{\pi_k \exp\big(-\lVert x_i - \mu_k \rVert^2 / 2s^2\big)}{\sum_j \pi_j \exp\big(-\lVert x_i - \mu_j \rVert^2 / 2s^2\big)} .$$

При $s \to 0$ слагаемое ближайшего среднего превосходит все остальные в неограниченно растущее число раз, поэтому $r_{ik} \to 1$ для ближайшего центра и $0$ для остальных: E-шаг превращается в шаг назначения метода k-средних. M-шаг для средних превращается в среднее назначенных точек — шаг обновления. А ожидаемый логарифм правдоподобия полных данных становится равным $-\frac{1}{2s^2} J + \text{const}$, где $J$ — целевая функция k-средних (Bishop, 2006, раздел 9.3.2). Алгоритм Ллойда — это EM для смеси одинаковых сферических гауссиан в пределе исчезающей дисперсии. В разделе «Код» этот EM запущен с $s = 0.01$: каждая ответственность равна 0 или 1, а средние совпадают с центрами k-средних из того же старта.

Предел объясняет и то, что предполагает k-средних, и то, что он теряет. Он предполагает кластеры с одинаковым сферическим разбросом. И он заменяет апостериорное распределение его наиболее вероятным значением — это жёсткий EM (hard EM), который уже не максимизирует правдоподобие. Точки в области перекрытия двух кластеров целиком отдаются одному из них, что раздвигает средние. На семи точках рисунка с нижней оценкой жёсткий EM, назначающий каждую точку ближайшему из $0$ и $\theta$, сходится к $\theta = 2.767$ против значения максимального правдоподобия $2.54$, найденного EM.

Локальные оптимумы, инициализация и вырожденные решения#

Монотонность гарантирует продвижение, а не пункт назначения. В разделе «Код» EM из 30 стартов со средними в случайных точках выборки приходит к шести разным значениям среднего логарифма правдоподобия, от $-3.703$ до $-3.259$, и 83% стартов достигают лучшего; при старте из центров k-means++ его достигают все 30. Отсюда обычный рецепт: инициализировать средние методом k-средних или k-means++ (scikit-learn по умолчанию сначала запускает k-средних), сделать несколько стартов и оставить тот, у которого правдоподобие наибольшее.

Начальные значения определяют и то, сколько времени займёт EM. В разделе «Код» нарочно плохой старт из трёх точек одного кластера держит средний логарифм правдоподобия около $-3.80$ с итерации 5 по итерацию 10, прежде чем он поднимается до $-3.259$: EM медленный, когда компоненты сильно перекрываются, потому что ответственности, а с ними и шаги, меняются мало.

Вырожденные решения требуют защиты. Компонента, захватившая одну-две точки, может сжимать свою ковариационную матрицу к вырожденной, пока правдоподобие неограниченно растёт. Три распространённых способа защиты: добавить небольшую константу к диагонали каждой $\Sigma_k$ (reg_covar в scikit-learn, по умолчанию $10^{-6}$); задать априорное распределение на ковариационных матрицах и запускать EM для MAP-оценки — регуляризованный M-шаг в смысле заметки о байесовском выводе; или перезапускать схлопывающуюся компоненту в другом месте. Ограничение ковариационных матриц — все равные, диагональные или сферические — уменьшает и риск, и число параметров.

Сколько компонент: BIC#

Правдоподобие растёт с $K$, поэтому само по себе не может выбрать $K$ — по той же причине, что и целевая функция k-средних. Штрафованное правдоподобие может. Байесовский информационный критерий (Bayesian information criterion, BIC; Schwarz, 1978) равен

$$\text{BIC} = -2\, \ell(\hat\theta) + p \log n ,$$

где $p$ — число свободных параметров; чем меньше, тем лучше. При больших $n$ он приближает $-2$, умноженное на логарифм маргинального правдоподобия модели, поэтому выбор наименьшего BIC приближает выбор модели с наибольшей апостериорной вероятностью при равных априорных вероятностях. У смеси с полными ковариационными матрицами в размерности $d$

$$p = (K - 1) + Kd + K\,\frac{d(d+1)}{2}$$

параметров: веса минус одно ограничение, средние и симметричные ковариационные матрицы. При $d = 2$ и $K = 3$ это $2 + 6 + 9 = 17$. На 300 точках из раздела «Код», порождённых тремя компонентами, BIC равен $2516.5, 2268.4, 2052.5, 2082.2, 2111.5, 2132.4$ для $K = 1, \dots, 6$: наименьший при $K = 3$. Альтернатива, не опирающаяся на асимптотику, — логарифм правдоподобия на отложенных данных: обучить модель на части данных и сравнить $K$ на остальных, как в заметке о валидации.

EM за пределами смесей#

В выводе нигде не использовались гауссианы. EM применим всякий раз, когда правдоподобие полных данных простое, а апостериорное распределение латентных переменных можно вычислить: пропущенные значения (латентные переменные — пропущенные элементы), смеси других распределений (смеси Бернулли для бинарных данных), скрытые марковские модели (алгоритм Баума–Велша, Baum et al., 1970, с марковской цепью состояний в роли латентной переменной) и восстановление истинных меток по нескольким зашумлённым разметчикам (Dawid and Skene, 1979). Следующая заметка снова встречает его в вероятностном PCA — модели с непрерывной латентной переменной. Когда апостериорное распределение нельзя вычислить точно, остаются два пути: сэмплировать из него сэмплером Гиббса из заметки о MCMC (EM методом Монте-Карло, Monte Carlo EM) или приблизить его лучшим $q$ из вычислимого семейства (вариационный EM, variational EM). Второй путь ведёт к VAE.

От EM к вариационному автоэнкодеру#

Для смеси E-шаг дешёв: апостериорное распределение каждой точки — это $K$ чисел, вычисленных по формуле Байеса. Теперь пусть латентная переменная — непрерывный вектор $z \in \mathbb R^{d'}$ с априорным распределением $\mathcal N(0, I)$, а декодер $p_\theta(x \mid z)$ — нейросеть. Апостериорному распределению $p_\theta(z \mid x) = p_\theta(x \mid z) p(z) / p_\theta(x)$ нужен интеграл $p_\theta(x) = \int p_\theta(x \mid z) p(z)\, dz$, у которого нет замкнутого вида; E-шаг невозможен. Вариационный автоэнкодер сохраняет нижнюю оценку и меняет способ её максимизации:

	EM для гауссовской смеси	Вариационный автоэнкодер
Латентная переменная	$z \in \{1, \dots, K\}$	$z \in \mathbb R^{d'}$, априорное распределение $\mathcal N(0, I)$
Модель $p_\theta(x \mid z)$	гауссиана для каждой компоненты	нейросетевой декодер
$q$	точное апостериорное распределение, $r_{ik}$ по формуле Байеса	энкодер $q_\varphi(z \mid x)$, гауссиана, среднее и дисперсию которой выдаёт сеть
E-шаг	точный, для каждой точки; зазор закрыт	один шаг градиента по $\varphi$, общий для всех точек (амортизированный, amortized); остаётся зазор $D_{\mathrm{KL}}(q_\varphi \,\|\, p_\theta(z \mid x))$
M-шаг	в замкнутом виде	один шаг градиента по $\theta$, совместно с $\varphi$
Математическое ожидание по $q$	сумма по $K$ значениям	оценка методом Монте-Карло с трюком репараметризации
Целевая функция	$\mathcal L(q, \theta) \le \log p_\theta(x)$	та же $\mathcal L$, называемая ELBO

Разложение, выведенное выше, $\log p_\theta(x) = \mathcal L + D_{\mathrm{KL}}(q \,\|\, p_\theta(z \mid x))$, — ровно то, которое используется в заметке о VAE, чтобы объяснить, насколько неточна ELBO. То, чего EM добивается точным E-шагом, VAE приближает энкодером, обученным на той же целевой функции, и цена в том, что максимизация нижней оценки лишь приближённо максимизирует правдоподобие. Диффузионные модели обучаются с вариационной границей того же рода, по целой цепочке латентных переменных (заметка о DDPM).

Что изменилось#

	k-средних	Гауссовская смесь с EM
Кластер	центр	вес, среднее и ковариационная матрица
Форма	сфера, одна для всех	любой эллипс, свой для каждого кластера
Принадлежность	жёсткая: $c_i$	мягкая: ответственности $r_{ik}$
Целевая функция	внутрикластерная сумма квадратов $J$	логарифм правдоподобия $\ell(\theta)$
Алгоритм	Ллойд: назначить, усреднить	EM: ответственности, взвешенная подгонка гауссиан
Гарантия	$J$ никогда не возрастает	$\ell$ никогда не убывает; предельные точки стационарны
Выбор $K$	локоть, силуэт	BIC, правдоподобие на отложенных данных
Выход	разбиение	плотность $p(x)$, пригодная для новых точек и для оценок аномальности

Ограничения#

• Гауссовские компоненты. Изогнутый кластер или кластер с тяжёлыми хвостами покрывается несколькими компонентами, и тогда кластеры смеси уже не совпадают с кластерами данных; двум лунам нужно много компонент.
• Локальные максимумы и медленная сходимость. EM находит стационарную точку, зависящую от старта, сходится линейно и медленно работает, когда компоненты перекрываются.
• Вырождение и число параметров. У каждой полной ковариационной матрицы $d(d+1)/2$ параметров; в высокой размерности они плохо оцениваются и появляются вырожденные решения, поэтому нужны диагональные или общие ковариационные матрицы, нижняя граница дисперсии или априорное распределение.
• Апостериорное распределение должно быть вычислимым. EM нужно точное апостериорное распределение латентных переменных; для большинства моделей с непрерывными латентными переменными и нелинейным правдоподобием оно недоступно, и E-шаг приходится приближать.
• BIC асимптотичен. Он предполагает много точек на параметр и регулярную модель; смеси нарушают условия регулярности, и на практике BIC — эвристика, которая часто работает.

Что нужно запомнить#

• Гауссовская смесь — это $p(x) = \sum_k \pi_k \mathcal N(x \mid \mu_k, \Sigma_k)$: выбрать компоненту, затем сэмплировать из её гауссианы; компонента — латентная переменная.
• Ответственности — апостериорные вероятности по формуле Байеса, $r_{ik} \propto \pi_k \mathcal N(x_i \mid \mu_k, \Sigma_k)$, нормированные по $k$.
• Для любого $q$ выполнено $\log p(x \mid \theta) = \mathcal L(q, \theta) + D_{\mathrm{KL}}(q \,\|\, p(z \mid x, \theta))$, где $\mathcal L = \mathbb E_q[\log p(x, z \mid \theta)] + H(q)$; нижняя оценка точна, когда $q$ — апостериорное распределение.
• EM — покоординатный подъём по $\mathcal L$: E-шаг делает $q$ апостериорным распределением, M-шаг максимизирует ожидаемый логарифм правдоподобия полных данных; логарифм правдоподобия никогда не убывает.
• M-шаг для гауссовской смеси: $\pi_k = N_k/n$, $\mu_k$ и $\Sigma_k$ — среднее и ковариация, взвешенные ответственностями.
• k-средних — это EM с равными сферическими ковариационными матрицами в пределе $s \to 0$, где ответственности становятся жёсткими назначениями.
• Инициализировать методом k-средних, делать несколько стартов, ограничивать дисперсии снизу и выбирать $K$ по BIC $= -2\ell + p \log n$ или по правдоподобию на отложенных данных.
• VAE максимизирует ту же нижнюю оценку с энкодером вместо точного апостериорного распределения; оставшийся зазор — KL-дивергенция энкодера относительно апостериорного распределения.

Код#

Ответственности, EM вручную, его совпадение со scikit-learn, предел k-средних, локальные оптимумы и BIC:

import numpy as np
from scipy.stats import norm, multivariate_normal as mvn
from scipy.special import logsumexp
from sklearn.mixture import GaussianMixture
from sklearn.cluster import KMeans

# ответственности по формуле Байеса: одномерная смесь 0.3 N(0, 1) + 0.7 N(4, 2^2) в точке x = 2
w = np.array([0.3, 0.7]) * norm.pdf(2, [0, 4], [1, 2])
print(w.round(5), round(w.sum(), 5), (w / w.sum()).round(4))     # [0.0162 0.08469] 0.10089 [0.1605 0.8395]

# 300 точек из смеси трёх компонент на плоскости
rng = np.random.default_rng(0)
mus = np.array([[-2.3, 1.2], [2.0, 2.0], [1.2, -1.8]])
Ss = np.array([[[1.0, 0.75], [0.75, 0.9]], [[0.25, 0], [0, 0.25]], [[1.6, -0.2], [-0.2, 0.5]]])
X = np.vstack([rng.multivariate_normal(m, S, size=k) for m, S, k in zip(mus, Ss, (120, 80, 100))])
n, d, K = len(X), 2, 3

def em(X, pi, mu, S, iters):
    lls = []
    for _ in range(iters):
        logw = np.log(pi) + np.stack([mvn(mu[k], S[k]).logpdf(X) for k in range(K)], axis=1)   # n x K
        lls.append(logsumexp(logw, axis=1).sum())                  # логарифм правдоподобия текущих параметров
        r = np.exp(logw - logsumexp(logw, axis=1, keepdims=True))  # E-шаг: ответственности
        Nk = r.sum(axis=0)                                          # M-шаг: взвешенные счётчики, средние, ковариации
        pi, mu = Nk / n, (r.T @ X) / Nk[:, None]
        S = np.stack([(r[:, k, None] * (X - mu[k])).T @ (X - mu[k]) / Nk[k] for k in range(K)])
    return pi, mu, S, np.array(lls)

pi0, mu0, S0 = np.full(K, 1 / K), X[[0, 1, 2]], np.stack([np.eye(d)] * K)   # три точки первой компоненты
pi, mu, S, lls = em(X, pi0, mu0, S0, 200)
print((lls[[0, 1, 5, 10, 25, 50, 199]] / n).round(4), bool(np.all(np.diff(lls) >= -1e-9)))
# [-8.6004 -4.0704 -3.8026 -3.7629 -3.3832 -3.2591 -3.2591] True: логарифм правдоподобия никогда не убывает
print(pi.round(3), mu.round(2).tolist())   # [0.337 0.401 0.263] [[1.29, -1.75], [-2.19, 1.36], [1.98, 1.98]]
gm = GaussianMixture(K, weights_init=pi0, means_init=mu0, precisions_init=S0, reg_covar=0, tol=1e-10, max_iter=500).fit(X)
print(round(gm.score(X), 4))               # -3.2591: scikit-learn из того же старта приходит к тому же максимуму

# мягкий и жёсткий EM на игрушечной задаче с рисунка нижней оценки: 0.5 N(0, 1) + 0.5 N(theta, 1), семь точек, theta0 = 0.6
x = np.array([-0.8, -0.2, 0.3, 0.9, 2.1, 2.8, 3.4])
soft = hard = 0.6
for _ in range(100):
    r = norm.pdf(x - soft) / (norm.pdf(x) + norm.pdf(x - soft))
    soft = (r * x).sum() / r.sum()                                 # среднее, взвешенное ответственностями
    hard = x[np.abs(x - hard) < np.abs(x)].mean()                  # среднее точек, которые ближе к theta, чем к 0
print(round(soft, 3), round(hard, 3))                              # 2.54 2.767

# k-средних как предел: ковариации зафиксированы, s^2 I с малым s, обновляются только средние
mu_h = X[[0, 150, 250]].copy()
for _ in range(50):
    logw = -((X[:, None, :] - mu_h[None]) ** 2).sum(axis=2) / (2 * 0.01 ** 2)
    r = np.exp(logw - logsumexp(logw, axis=1, keepdims=True))
    mu_h = (r.T @ X) / r.sum(axis=0)[:, None]
km = KMeans(K, init=X[[0, 150, 250]], n_init=1).fit(X)
print(r.max(axis=1).min().round(6), np.abs(mu_h - km.cluster_centers_).max().round(10))   # 1.0 0.0: жёсткие назначения, центры k-средних

# локальные оптимумы: 30 стартов из случайных точек выборки и 30 из центров k-means++
for init in ('random_from_data', 'k-means++'):
    ll = np.array([GaussianMixture(K, init_params=init, max_iter=2000, tol=1e-8, random_state=s).fit(X).score(X) for s in range(30)])
    print(init, np.unique(ll.round(3)), round((ll > ll.max() - 1e-3).mean(), 3))
# random_from_data [-3.703 -3.701 -3.697 -3.656 -3.636 -3.259] 0.833 | k-means++ [-3.259] 1.0

# выбор K по BIC = -2 log L + p log n, p = (K - 1) + K d + K d (d + 1) / 2
for k in range(1, 7):
    g = GaussianMixture(k, n_init=5, random_state=0).fit(X)
    p = (k - 1) + k * d + k * d * (d + 1) // 2
    print(k, p, round(g.bic(X), 1), round(-2 * g.score(X) * n + p * np.log(n), 1))
# 1 5 2516.5 2516.5 | 2 11 2268.4 2268.4 | 3 17 2052.5 2052.5 | 4 23 2082.2 2082.2 | 5 29 2111.5 2111.5 | 6 35 2132.4 2132.4

Дальше#

Смесь объясняет каждую точку дискретной латентной переменной: из какого из $K$ кластеров она пришла. Многие наборы данных лучше объясняются непрерывной: изображение цифры из $64$ пикселей меняется вдоль нескольких направлений (наклон, ширина, толщина штриха), и данные лежат вблизи подпространства или поверхности малой размерности. Следующая заметка находит такие направления — сначала линейно, методом главных компонент, затем с помощью его вероятностной версии, которая снова оказывается моделью с латентными переменными, обучаемой методом максимального правдоподобия, и наконец нелинейными отображениями для визуализации.

Литература#

• A. P. Dempster, N. M. Laird, D. B. Rubin. Maximum Likelihood from Incomplete Data via the EM Algorithm. Journal of the Royal Statistical Society B, 39(1), 1–38, 1977.
• C. F. J. Wu. On the Convergence Properties of the EM Algorithm. The Annals of Statistics, 11(1), 95–103, 1983.
• R. M. Neal, G. E. Hinton. A View of the EM Algorithm that Justifies Incremental, Sparse, and other Variants. In Learning in Graphical Models, 355–368, 1998.
• G. Schwarz. Estimating the Dimension of a Model. The Annals of Statistics, 6(2), 461–464, 1978.
• L. E. Baum, T. Petrie, G. Soules, N. Weiss. A Maximization Technique Occurring in the Statistical Analysis of Probabilistic Functions of Markov Chains. The Annals of Mathematical Statistics, 41(1), 164–171, 1970.
• A. P. Dawid, A. M. Skene. Maximum Likelihood Estimation of Observer Error-Rates Using the EM Algorithm. Applied Statistics, 28(1), 20–28, 1979.
• C. M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006. Разделы 9.2 (смеси гауссиан), 9.3 (альтернативный взгляд на EM; 9.3.2 — связь с k-средних), 9.4 (EM в общем виде).
• K. P. Murphy. Probabilistic Machine Learning: An Introduction. MIT Press, 2022. Раздел 8.7 (оптимизация через оценки и EM).
• D. P. Kingma, M. Welling. Auto-Encoding Variational Bayes. ICLR, 2014.
• E. Sokolov et al. Машинное обучение (HSE), конспекты лекций, лекция 15: EM-алгоритм (2025).