Mixtures and the EM algorithm

The clustering note ended with methods that give hard answers. k-means puts each point in exactly one cluster, says nothing about how sure it is, and behaves as if all clusters were round and equally wide. This note turns clustering into a probability model. Each cluster becomes a Gaussian with its own mean, covariance and weight, the data are a sample from their mixture, and the parameters are fitted by maximum likelihood. The likelihood has no closed-form maximum, because it contains the cluster memberships as unobserved variables. The EM algorithm handles such models by climbing a lower bound on the log-likelihood, one block of variables at a time. The same bound, with the exact posterior replaced by a neural network, is the training objective of the variational autoencoder.

Notation. $x_1, \dots, x_n \in \mathbb R^d$ are the data, $K$ the number of components. $z_i \in \{1, \dots, K\}$ is the unobserved component of point $i$. A component has a weight $\pi_k \ge 0$ with $\sum_k \pi_k = 1$, a mean $\mu_k$ and a covariance $\Sigma_k$, symmetric positive definite; $\theta$ collects all of them. $\mathcal N(x \mid \mu, \Sigma)$ is the Gaussian density with mean $\mu$ and covariance $\Sigma$ (variance, not standard deviation, as the second argument). $\ell(\theta) = \sum_i \log p(x_i \mid \theta)$ is the log-likelihood, $r_{ik}$ the responsibilities defined below, $q$ a distribution over the latent variables.

The Gaussian mixture model#

A Gaussian mixture describes the data by a two-step generative story. First pick a component $z$ with probabilities $\pi_1, \dots, \pi_K$; then draw $x$ from that component’s Gaussian:

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

The component is never observed, so the distribution of $x$ alone is the marginal, obtained by summing the joint $p(x, z = k) = \pi_k \mathcal N(x \mid \mu_k, \Sigma_k)$ over $k$ with the law of total probability:

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

Each component is a multivariate Gaussian, a cloud whose contours are ellipses set by the eigenvectors of $\Sigma_k$, so a mixture can describe clusters that are elongated, tilted, small or wide, and with enough components it approximates any smooth density. Compared with k-means, the mixture has three new kinds of parameters: a covariance per cluster, a weight per cluster, and a probability, rather than a verdict, for each point’s membership.

Worked example. In one dimension, take $\pi = (0.3, 0.7)$, $\mu = (0, 4)$ and standard deviations $1$ and $2$. At $x = 2$ the two component densities are $\mathcal N(2 \mid 0, 1) = 0.0540$ and $\mathcal N(2 \mid 4, 4) = 0.1210$, so

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

The likelihood with latent variables#

For i.i.d. data the log-likelihood is

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

The logarithm of a sum does not split. Setting the gradient to zero couples every parameter with every other through the sum in the denominator, and there is no closed-form solution. Compare the complete-data log-likelihood, the one we would have if the components $z_i$ were observed:

$$\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) .$$

This is a sum of logarithms. It separates into one Gaussian fit per component, on the points labelled with that component, and one categorical fit for the weights; each has the closed form of the maximum likelihood note: the share of points, the mean and the $1/n_k$ covariance. The difficulty is entirely in not knowing $z$. The EM algorithm turns that observation into a method: replace the unknown labels by their probabilities given the data, fit as if the labels were known, and repeat.

Three properties of $\ell$ are worth knowing before fitting it. It is not concave and has several local maxima. It is symmetric under relabelling the components, so every maximum comes in $K!$ copies. And it is unbounded: a component centred on a single data point whose variance shrinks to zero makes $\ell \to \infty$, the case worked out in the maximum likelihood note. The unrestricted maximum likelihood estimate of a Gaussian mixture therefore does not exist: the supremum is infinite and no parameter attains it. What EM returns is a local maximum of the likelihood away from the collapsing solutions or, with a covariance floor or a prior as in the section on degenerate solutions below, the maximizer of a constrained or regularized likelihood.

Responsibilities#

Given the parameters, the probability that point $x_i$ came from component $k$ follows from Bayes’ formula: prior $\pi_k$, likelihood $\mathcal N(x_i \mid \mu_k, \Sigma_k)$, evidence $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)} .$$

The $r_{ik}$ are the responsibilities: each row is a distribution over components, $r_{ik} \ge 0$ and $\sum_k r_{ik} = 1$. They are the soft version of the k-means assignment. In the worked example above, the point $x = 2$ belongs to the first component with probability $0.0162/0.1009 = 0.16$ and to the second with $0.84$: it lies two standard deviations from the first mean and one from the second, and the second component is also more frequent. k-means, comparing distances only, would put $x = 2$ halfway, on the boundary.

Where the lower bound comes from#

Take one point $x$ with latent $z$ (for the whole dataset, sum over points), and let $q(z)$ be any distribution over the $K$ values of $z$ that is positive wherever the posterior is. Since $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)}$$

for every $z$. The left side does not depend on $z$, so averaging both sides over $z \sim q$ leaves it unchanged:

$$\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))} .$$

The second term is the Kullback–Leibler divergence from $q$ to the posterior. It is non-negative and zero exactly when $q$ equals the posterior. Hence

$$\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),$$

with equality when $q$ is the posterior $p(z \mid x, \theta)$. $\mathcal L$ is the evidence lower bound (ELBO); Neal and Hinton (1998) call its negative the free energy. The same inequality follows in one line from Jensen’s inequality for the concave logarithm, $\log \mathbb E_q[p(x, z)/q(z)] \ge \mathbb E_q[\log p(x, z)/q(z)]$, but the identity above says more: the gap is exactly a KL divergence, so we know when the bound is tight.

The bound is useful because its first term is the complete-data log-likelihood averaged over guessed labels, a sum of logarithms that we can maximize in closed form, and its second term does not depend on $\theta$.

EM is coordinate ascent on the bound#

$\mathcal L(q, \theta)$ depends on two blocks, the distribution $q$ (for i.i.d. data, one $q_i$ per point) and the parameters $\theta$. The EM algorithm (Dempster, Laird and Rubin, 1977) maximizes it over one block at a time:

1. E-step. With $\theta^t$ fixed, maximize over $q$. The left side $\log p(x \mid \theta^t)$ does not depend on $q$, so maximizing $\mathcal L$ means minimizing the KL term, and the maximum is the posterior: $q^{t+1}(z) = p(z \mid x, \theta^t)$. For a mixture this is computing the responsibilities $r_{ik}$. After the E-step the bound touches the log-likelihood at $\theta^t$.
2. M-step. With $q^{t+1}$ fixed, maximize over $\theta$. The entropy $H(q^{t+1})$ is a constant, so

$$\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] ,$$

the expected complete-data log-likelihood, called $Q(\theta, \theta^t)$ in the original paper.

Monotonicity follows in one line, using the bound, then the M-step, then the tightness after the E-step:

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

The log-likelihood never decreases. Any $\theta$ that increases the bound, not only its maximizer, keeps the inequality, which is why an M-step can be replaced by a few gradient steps (generalized EM). Two more facts complete the picture. At $\theta^t$ the bound $\mathcal L(q^{t+1}, \cdot)$ and $\ell$ touch and have the same gradient, because their difference, a KL divergence, has its minimum 0 there; so EM stops moving only at stationary points of $\ell$. Never decreasing is not yet converging. $\ell(\theta^t)$ converges to a finite value when $\ell$ is bounded above along the iterates, and the parameters need more. Wu (1983) shows that if the set $\{\theta : \ell(\theta) \ge \ell(\theta^0)\}$ is compact and $\ell$ and the bound are continuous, every limit point of the EM iterates is a stationary point of $\ell$ and $\ell(\theta^t)$ converges monotonically; a stationary point can be a local maximum or a saddle, never guaranteed to be the global maximum. For a Gaussian mixture compactness fails at the collapsing solutions above, so the guarantee covers only runs that stay away from them, for instance with a floor on the covariances.

Figure

EM climbs a lower bound

Worked example. The figure’s model has a single unknown parameter, the mean $\theta$ of the second of two equally weighted unit-variance components, the first fixed at $\mathcal N(0, 1)$. At $\theta_0 = 0.6$ the responsibilities of the second component for the seven points $-0.8, -0.2, 0.3, 0.9, 2.1, 2.8, 3.4$ are $0.341, 0.426, 0.500, 0.589, 0.746, 0.818, 0.865$: points far to the right are attributed mostly to it. With the weights and the first component fixed, the M-step for $\theta$ is a weighted mean,

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

The log-likelihood rises from $-16.35$ to $-12.83$, then to $-12.14$ at $\theta_2 = 2.168$, and EM converges to the maximum at $\theta = 2.54$ with $\ell = -12.00$. The steps shrink as the bound and the log-likelihood become alike near the maximum: EM converges linearly, fast at first and slowly at the end.

The updates for Gaussian mixtures#

For a Gaussian mixture the expected complete-data log-likelihood is

$$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),$$

the complete-data log-likelihood with each indicator $[z_i = k]$ replaced by its posterior probability $r_{ik}$. It separates by parameter. Write $N_k = \sum_i r_{ik}$ for the effective number of points of component $k$.

Weights. Maximize $\sum_k N_k \log \pi_k$ subject to $\sum_k \pi_k = 1$. The Lagrangian $\sum_k N_k \log \pi_k + \lambda(\sum_k \pi_k - 1)$ has derivative $N_k/\pi_k + \lambda = 0$, so $\pi_k \propto N_k$, and normalizing, $\pi_k = N_k / n$: the share of points, with fractional counts.

Means. The terms with $\mu_k$ are $-\frac12 \sum_i r_{ik} (x_i - \mu_k)^\top \Sigma_k^{-1} (x_i - \mu_k)$. The gradient $\Sigma_k^{-1} \sum_i r_{ik} (x_i - \mu_k)$ vanishes when $\sum_i r_{ik}(x_i - \mu_k) = 0$, because $\Sigma_k^{-1}$ is invertible.

Covariances. With the precision $\Lambda_k = \Sigma_k^{-1}$, the terms with $\Sigma_k$ are $\frac{N_k}{2} \log \det \Lambda_k - \frac12 \operatorname{tr}(\Lambda_k S_k)$, where $S_k = \sum_i r_{ik} (x_i - \mu_k)(x_i - \mu_k)^\top$. Using $\partial \log\det \Lambda / \partial \Lambda = \Lambda^{-1}$ and $\partial \operatorname{tr}(\Lambda S)/\partial \Lambda = S$, the gradient $\frac{N_k}{2}\Lambda_k^{-1} - \frac12 S_k$ vanishes at $\Sigma_k = S_k / N_k$.

Together, the M-step is

$$\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 ,$$

the Gaussian maximum likelihood estimates with each point counted with weight $r_{ik}$ (the new $\mu_k$ is used in $\Sigma_k$). One EM iteration costs $O(nKd^2)$ for the responsibilities plus $O(Kd^3)$ for the inverses and determinants. In practice the E-step is computed in logarithms, with $\log \sum_k \exp(\cdot)$ evaluated stably, because the densities underflow in high dimension.

Figure

EM for a mixture of three Gaussians

The figure starts EM badly on purpose, with all three means inside one cloud. After one iteration the components have already spread out, after five the elongated and the small clouds are found, and by iteration 25 the third component has settled on the wide cloud. The average log-likelihood climbs from $-7.020$ to $-3.223$ and never falls. The ellipses show what k-means could not represent: one tilted, one small and one wide cluster, each with its own covariance.

k-means is the hard, equal-variance limit#

Fix every covariance at $\Sigma_k = s^2 I$ with a known $s$, and update only the means (the weights may be fixed or updated; the limit does not depend on them). The responsibilities become

$$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)} .$$

As $s \to 0$ the term of the nearest mean dominates every other by a factor that grows without bound, so $r_{ik} \to 1$ for the nearest centre and $0$ otherwise: the E-step becomes the assignment step of k-means. The M-step for the means becomes the mean of the assigned points: the update step. And the expected complete-data log-likelihood becomes $-\frac{1}{2s^2} J + \text{const}$, with $J$ the k-means objective (Bishop, 2006, Section 9.3.2). Lloyd’s algorithm is EM for a mixture of identical spherical Gaussians in the limit of vanishing variance. The Code section runs this EM with $s = 0.01$: every responsibility is 0 or 1 and the means coincide with the k-means centres from the same start.

The limit explains both what k-means assumes and what it loses. It assumes clusters with equal spherical spread. And it replaces the posterior by its most probable value, hard EM, which no longer maximizes the likelihood. Points in the overlap of two clusters are given entirely to one of them, which pulls the means apart. On the seven points of the bound figure, hard EM, which assigns each point to the nearer of $0$ and $\theta$, converges to $\theta = 2.767$, against the maximum likelihood value $2.54$ found by EM.

Local optima, initialization and degenerate solutions#

Monotonicity guarantees progress, not the destination. In the Code section, EM from 30 starts with means at random data points ends at six different values of the average log-likelihood, between $-3.703$ and $-3.259$, and 83% of the starts reach the best one; started from k-means++ centres, all 30 reach it. The usual recipe follows: initialize the means with k-means or k-means++ (scikit-learn runs k-means first by default), run several starts, and keep the one with the highest likelihood.

Starting values also decide how long EM takes. In the Code section a deliberately poor start from three points of one cluster leaves the average log-likelihood near $-3.80$ from iteration 5 to iteration 10 before it climbs to $-3.259$: EM is slow when components overlap heavily, because the responsibilities, and with them the steps, change little.

Degenerate solutions need protection. A component that captures one or two points can shrink its covariance toward a singular matrix while the likelihood grows without bound. Three common defences: add a small constant to the diagonal of every $\Sigma_k$ (scikit-learn’s reg_covar, $10^{-6}$ by default); put a prior on the covariances and run EM for the maximum a posteriori estimate, a regularized M-step in the sense of the Bayesian inference note; or restart a collapsing component elsewhere. Constraining the covariances, all equal, diagonal or spherical, reduces both the risk and the number of parameters.

How many components: BIC#

The likelihood grows with $K$, so it cannot choose $K$ by itself, for the same reason that the k-means objective cannot. A penalized likelihood can. The Bayesian information criterion (Schwarz, 1978) is

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

where $p$ is the number of free parameters; smaller is better. It approximates $-2$ times the log of the marginal likelihood of the model for large $n$, so choosing the smallest BIC approximates choosing the model with the highest posterior probability under equal prior probabilities. A full-covariance mixture in $d$ dimensions has

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

parameters: the weights minus one constraint, the means, and the symmetric covariances. For $d = 2$ and $K = 3$ that is $2 + 6 + 9 = 17$. On the 300 points of the Code section, drawn from three components, the BIC is $2516.5, 2268.4, 2052.5, 2082.2, 2111.5, 2132.4$ for $K = 1, \dots, 6$: smallest at $K = 3$. Held-out log-likelihood is the alternative that does not rely on asymptotics: fit on part of the data and compare $K$ on the rest, as in the validation note.

EM beyond mixtures#

Nothing in the derivation used Gaussians. EM applies whenever the complete-data likelihood is easy and the posterior over the latent variables can be computed: missing values (the latent variables are the missing entries), mixtures of other distributions (Bernoulli mixtures for binary data), hidden Markov models (the Baum–Welch algorithm, Baum et al., 1970, with the Markov chain of states as the latent variable), and the reconstruction of true labels from several noisy annotators (Dawid and Skene, 1979). The next note meets it again in probabilistic PCA, a model with a continuous latent variable. When the posterior cannot be computed exactly, two routes remain: sample it, with the Gibbs sampler of the MCMC note (Monte Carlo EM), or approximate it with the best $q$ from a tractable family (variational EM). The second route leads to the VAE.

From EM to the variational autoencoder#

For a mixture, the E-step is cheap: the posterior of each point is $K$ numbers, computed by Bayes’ formula. Now let the latent variable be a continuous vector $z \in \mathbb R^{d'}$ with prior $\mathcal N(0, I)$, and let the decoder $p_\theta(x \mid z)$ be a neural network. The posterior $p_\theta(z \mid x) = p_\theta(x \mid z) p(z) / p_\theta(x)$ needs the integral $p_\theta(x) = \int p_\theta(x \mid z) p(z)\, dz$, which has no closed form; the E-step is impossible. The variational autoencoder keeps the bound and changes how it is maximized:

	EM for a Gaussian mixture	Variational autoencoder
Latent variable	$z \in \{1, \dots, K\}$	$z \in \mathbb R^{d'}$, prior $\mathcal N(0, I)$
Model $p_\theta(x \mid z)$	a Gaussian per component	a neural decoder
$q$	the exact posterior, $r_{ik}$ by Bayes’ formula	an encoder $q_\varphi(z \mid x)$, a Gaussian whose mean and variance a network outputs
E-step	exact, per point; gap closed	one gradient step on $\varphi$, shared by all points (amortized); gap $D_{\mathrm{KL}}(q_\varphi \,\|\, p_\theta(z \mid x))$ remains
M-step	closed form	one gradient step on $\theta$, jointly with $\varphi$
Expectation over $q$	a sum over $K$ values	a Monte Carlo estimate with the reparameterization trick
Objective	$\mathcal L(q, \theta) \le \log p_\theta(x)$	the same $\mathcal L$, called the ELBO

The decomposition derived above, $\log p_\theta(x) = \mathcal L + D_{\mathrm{KL}}(q \,\|\, p_\theta(z \mid x))$, is exactly the one used in the VAE note to explain how loose the ELBO is. What EM achieves by an exact E-step, the VAE approximates with an encoder trained on the same objective, and the price is that maximizing the bound only approximately maximizes the likelihood. Diffusion models are trained with a bound of the same kind, over a whole chain of latent variables (DDPM note).

What changed#

	k-means	Gaussian mixture with EM
Cluster	a centre	a weight, a mean and a covariance
Shape	sphere, same for all	any ellipse, its own for each cluster
Membership	hard: $c_i$	soft: responsibilities $r_{ik}$
Objective	within-cluster sum of squares $J$	log-likelihood $\ell(\theta)$
Algorithm	Lloyd: assign, average	EM: responsibilities, weighted Gaussian fits
Guarantee	$J$ never increases	$\ell$ never decreases; limit points are stationary
Choosing $K$	elbow, silhouette	BIC, held-out likelihood
Output	a partition	a density $p(x)$, usable for new points and for anomaly scores

Limitations#

• Gaussian components. A curved or heavy-tailed cluster is covered by several components, and the clusters of the mixture then no longer match the clusters of the data; the two moons need many components.
• Local maxima and slow convergence. EM finds a stationary point that depends on the start, converges linearly, and is slow when components overlap.
• Degeneracy and parameter count. Full covariances have $d(d+1)/2$ parameters each; in high dimension they are poorly estimated and singular solutions appear, so diagonal or tied covariances, a variance floor or a prior are needed.
• The posterior must be computable. EM needs the exact posterior over the latent variables; for most models with continuous latent variables and nonlinear likelihoods it is not available, and the E-step must be approximated.
• BIC is asymptotic. It assumes many points per parameter and a regular model; mixtures violate the regularity conditions, and in practice BIC is a heuristic that often works.

What you should remember#

• A Gaussian mixture is $p(x) = \sum_k \pi_k \mathcal N(x \mid \mu_k, \Sigma_k)$: pick a component, then draw from its Gaussian; the component is a latent variable.
• Responsibilities are posteriors by Bayes’ formula, $r_{ik} \propto \pi_k \mathcal N(x_i \mid \mu_k, \Sigma_k)$, normalized over $k$.
• For any $q$, $\log p(x \mid \theta) = \mathcal L(q, \theta) + D_{\mathrm{KL}}(q \,\|\, p(z \mid x, \theta))$ with $\mathcal L = \mathbb E_q[\log p(x, z \mid \theta)] + H(q)$; the bound is tight when $q$ is the posterior.
• EM is coordinate ascent on $\mathcal L$: the E-step sets $q$ to the posterior, the M-step maximizes the expected complete-data log-likelihood; the log-likelihood never decreases.
• Gaussian mixture M-step: $\pi_k = N_k/n$, $\mu_k$ and $\Sigma_k$ the responsibility-weighted mean and covariance.
• k-means is EM with equal spherical covariances in the limit $s \to 0$, where responsibilities become hard assignments.
• Initialize with k-means, use several starts, floor the variances, and choose $K$ by BIC $= -2\ell + p \log n$ or held-out likelihood.
• The VAE maximizes the same bound with an encoder in place of the exact posterior; the remaining gap is the KL from the encoder to the posterior.

Code#

Responsibilities, EM by hand, its agreement with scikit-learn, the k-means limit, local optima and 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

# responsibilities by Bayes' formula: a 1-D mixture 0.3 N(0, 1) + 0.7 N(4, 2^2) at 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 points from a three-component mixture in 2-D
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())                  # log-likelihood of the current parameters
        r = np.exp(logw - logsumexp(logw, axis=1, keepdims=True))  # E-step: responsibilities
        Nk = r.sum(axis=0)                                          # M-step: weighted counts, means, covariances
        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)   # three points of the first component
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: the log-likelihood never decreases
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 reaches the same maximum from the same start

# soft and hard EM on the toy of the bound figure: 0.5 N(0, 1) + 0.5 N(theta, 1), seven points, 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()                                 # responsibility-weighted mean
    hard = x[np.abs(x - hard) < np.abs(x)].mean()                  # mean of the points nearer to theta than to 0
print(round(soft, 3), round(hard, 3))                              # 2.54 2.767

# k-means as the limit: covariances fixed at s^2 I with a small s, only the means updated
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: hard assignments, the k-means centres

# local optima: 30 starts from random data points, and 30 from k-means++ centres
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

# choosing K with 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

Next#

A mixture explains each point by a discrete latent variable: which of $K$ clusters it came from. Many datasets are better explained by a continuous one: a $64$-pixel image of a digit varies along a few directions (slant, width, stroke thickness), and the data lie near a low-dimensional subspace or surface. The next note finds such directions, first linearly with principal component analysis, then with its probabilistic version, which is again a latent-variable model fitted by maximum likelihood, and finally with nonlinear maps for visualization.

References#

• 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. Sections 9.2 (mixtures of Gaussians), 9.3 (an alternative view of EM; 9.3.2, the relation to k-means), 9.4 (EM in general).
• K. P. Murphy. Probabilistic Machine Learning: An Introduction. MIT Press, 2022. Section 8.7 (bound optimization and EM).
• D. P. Kingma, M. Welling. Auto-Encoding Variational Bayes. ICLR, 2014.
• E. Sokolov et al. Машинное обучение (HSE), lecture notes, lecture 15: the EM algorithm (2025).