Expectation-Maximization as Inference with Latent Variables

Latent variables make a likelihood awkward for one structural reason. With observed data $x$ and hidden variables $z$ the likelihood is

$$\mathcal{L}(\theta) = \log p(x \mid \theta) = \log \sum_{z} p(z \mid \theta),p(x \mid z, \theta)$$

The logarithm sits outside a sum, so terms that factor cleanly inside the sum do not factor at all outside it. Maximizing the complete-data log likelihood $\log p(x, z \mid \theta)$ would be easy for most models; the difficulty is that $z$ was never observed, and averaging over it inside a logarithm couples the parameters of every mixture component to the parameters of every other. Expectation maximization resolves the coupling by a different route from the one sampling takes: instead of drawing states of $z$, it averages over them with weights it recomputes each iteration, and the result is a sequence of parameter values whose likelihood never decreases.

Maximizing the expected complete-data log likelihood

Suppose the current parameter estimate is $\theta_t$. Averaging the complete-data log likelihood over the posterior of the latents under that estimate gives a function of the new parameters alone:

$$Q(\theta \mid \theta_t) = \mathbb{E}_{p(z \mid x, \theta_t)}\bigl[\log p(x, z \mid \theta)\bigr] = \sum_z p(z \mid x, \theta_t) \log p(x, z \mid \theta)$$

The E step is the computation of the posterior weights $p(z \mid x, \theta_t)$; the M step is $\theta_{t+1} = \arg\max_\theta Q(\theta \mid \theta_t)$. The idea is to pick $\theta$ that scores better on the completed data than $\theta_t$ does, and then to recompute the weights, since they depended on $\theta_t$.

The reason the loop does not diverge is worth spelling out, since it is the argument that Variational Inference and the Evidence Lower Bound generalizes. For any distribution $q$ over $z$, $\log p(x \mid \theta) = \mathrm{KL}(q | p(z \mid x, \theta)) + \mathbb{E}_q[\log p(x,z\mid\theta) - \log q(z)]$ holds exactly, with the bound tight when $q$ equals the posterior. Setting $q = p(z \mid x, \theta_t)$ makes the bound tight at $\theta_t$, so the bound equals $\log p(x \mid \theta_t)$ there. The M step then maximizes a function that is a lower bound on the log likelihood everywhere and agrees with it at the previous iterate, so the new likelihood is at least the old one. Because $Q$ separates into $\sum_z q(z)\log p(x,z\mid\theta)$ and the term $-H(q)$ that does not involve $\theta$, maximizing the bound is the same as maximizing $Q$. Monotone non-decreasing likelihood, one tight bound per iteration, no tuning parameter.

A closed-form M step: the Gaussian mixture

The canonical case has one discrete latent per observation: observation $i$ belongs to component $k$ with unknown probability $\pi_k$, and given the component its value is normal with mean $\mu_k$ and variance $\sigma_k^2$. The E step computes responsibilities, the posterior probability that observation $i$ came from component $k$:

$$r_{ik} = \frac{\pi_k,\mathcal{N}(x_i \mid \mu_k, \sigma_k^2)} {\sum_{j} \pi_j,\mathcal{N}(x_i \mid \mu_j, \sigma_j^2)}$$

The M step is then weighted averaging, with the responsibilities as weights: each mean is the responsibility-weighted average of the data, each variance the responsibility-weighted average squared deviation, and each weight the mean responsibility. Every quantity is a closed-form function of the responsibilities, so one iteration costs one pass over the observations against the components, with a density evaluation per pair.

import math

def em(data, k=2, iters=50):
    n = len(data)
    span = max(data) - min(data)
    means = [min(data) + (j + 1) * span / (k + 1) for j in range(k)]
    var = [sum((x - sum(data) / n) ** 2 for x in data) / n] * k
    weight = [1.0 / k] * k

    def dens(x, j):
        return math.exp(-0.5 * (x - means[j]) ** 2 / var[j]) / math.sqrt(2 * math.pi * var[j])

    for _ in range(iters):
        resp = [[weight[j] * dens(x, j) for j in range(k)] for x in data]
        for row in resp:                       # responsibilities sum to one per observation
            s = sum(row)
            for j in range(k):
                row[j] /= s
        for j in range(k):
            nk = sum(row[j] for row in resp)
            if nk < 1e-12:                     # component lost every point; leave it untouched
                continue
            means[j] = sum(row[j] * x for row, x in zip(resp, data)) / nk
            var[j] = sum(row[j] * (x - means[j]) ** 2 for row, x in zip(resp, data)) / nk
            weight[j] = nk / n
    return weight, means, var

Fixed points are not necessarily maxima

The monotonicity guarantee is exactly that the likelihood does not decrease; it says nothing about where the iteration stops. Convergence is to a stationary point, which may be a local maximum, and the objective is generally not concave, so different starting values reach different solutions. In a mixture the likelihood is also invariant under permutation of the component labels, so the parameters that are found depend on the ordering assigned by the initialization. The standard mitigation is multiple restarts with the largest final likelihood retained, and the likelihood itself must be accumulated separately from the responsibilities, since monotonicity is a property of the sequence rather than a guarantee about any single step’s size.

A second class of fixed points is worse than a local maximum. If a component’s mean moves onto a single observation and its variance shrinks towards zero, that observation’s density under the component grows without bound, the likelihood rises towards infinity, and the responsibility of the data point converges to one. The M step then keeps the variance shrinking. The degenerate solution has infinite likelihood and describes nothing: it assigns one component per data point. Guarding against it requires either a floor on the variance parameter, a prior on the parameters that penalizes small variances, or a check that discards solutions with a component supported on fewer than a handful of points. The singularity is a property of the unbounded likelihood, not of the implementation.

Where the exact E step is unavailable

The E step requires the posterior $p(z \mid x, \theta_t)$ as a normalized distribution over all assignments of the latents. In a model with many coupled latents, such as a network with a wide cluster or a latent variable that several observations share, computing that posterior is exactly the inference problem of Exact Inference by Variable Elimination and the Cost of Treewidth, and it inherits the exponential cost in treewidth. When the graph is too wide the E step is replaced by the best member of a restricted family, which produces variational EM: monotonicity of the bound survives, because each step still improves the same lower bound, but the bound is no longer tight at the current iterate, so the fixed point is a maximum of the bound rather than of the likelihood, and the distance between the two can be substantial. The alternative is to replace the E step with samples from the posterior as in Metropolis-Hastings: Detailed Balance and Proposal Design, which makes the average over $z$ a Monte Carlo estimate; the likelihood is then no longer monotone in general, since the estimate is noisy.

Averaging, sampling, or maximizing the hidden variables

The E step is one of three ways to handle the latents, and the choices differ in which guarantee they keep. Averaging with the exact posterior, which is what EM does, preserves the monotone increase in the likelihood and costs one inference pass per iteration; it is the only one of the three that cannot decrease the objective. Sampling the latents from the posterior instead makes the E step unbiased for the same quantity but destroys monotonicity, because each iteration optimizes a different noisy version of the objective; the likelihood then jitters rather than climbs, and the number of samples per iteration controls the size of the jitter rather than its presence.

Maximizing over the latents instead of averaging them defines a different iteration. It is still a generalized EM step in the formal sense — the E step maximizes the bound within the restricted family of point-mass distributions over assignments — but that bound is lower than the one attained by the posterior, so the fixed points of the two versions need not coincide. In a mixture model the hard-assignment limit is approached as the component covariances shrink relative to the separation of their means, since the responsibilities become one-hot in that limit, which is why soft and hard versions agree on well-separated data and disagree where responsibilities are spread across components. The same distinction separates the exact E step from the variational E step: a substitution of expected statistics inside the logarithm is not the expectation of the logarithm, and that inequality is where the approximation error enters.

Convergence speed and identifiability

Iterations are cheap relative to the alternatives, but their number is not bounded by anything in the model. EM converges linearly, and the rate degrades towards one as the fraction of missing information grows, so a model in which the latents are nearly determined by the parameters reaches a fixed point quickly while a model in which the latents remain uncertain takes many more iterations for the same improvement. Augmented variants that take a step past the M-step optimum, or that use a Newton direction on the observed-data likelihood, trade per-iteration cost for fewer iterations.

Identifiability is a separate matter from convergence. Parameters that make the observed likelihood identical for all values cannot be distinguished by this algorithm or any other: component labels in a mixture, scale and location in some factor models, and the orientation of latent directions in others. The practical consequence for the fixed point is that the reported parameters are one representative of an equivalence class selected by initialization, and that a change of the initialization is the only cheap way to see whether the solution was determined by the data.