Learning with Hidden Variables: The EM Algorithm

This builds on Learning Probabilistic Models, which framed learning as Bayesian inference and derived the maximum-likelihood, MAP, and Bayesian estimators for the case of complete data — where every example fixes a value for every variable, so each parameter is just a count. Here we drop that assumption. When some variables are never observed, counting is impossible, and the learning problem needs the algorithm this part is about.

Learning with hidden variables: the EM algorithm

Everything so far assumed complete data. Real data are usually incomplete: some variables are never observed. A medical record lists symptoms, diagnosis, and treatment, but rarely the underlying disease. Such a hidden (or latent) variable is worth keeping precisely because it can dramatically reduce the number of parameters: a disease node with three predisposing causes and three symptoms may need $78$ parameters, but splicing the disease out and connecting causes directly to symptoms — now no longer conditionally independent — can balloon that to over $700$. Hidden variables buy compactness at the cost of a harder learning problem: you cannot count what you cannot see.

The expectation–maximization algorithm, EM, solves this in a general way. Its one idea: pretend you know the parameters, use them to infer a probability distribution over the hidden variables, then refit the parameters as if that inferred completion were real data — and iterate. Each pass alternates an expectation step (complete the data in expectation) with a maximization step (re-estimate the parameters).

Unsupervised clustering with a mixture of Gaussians

The cleanest instance is unsupervised clustering: discovering categories in unlabeled data. Model the data as generated by a mixture distribution with $k$ components. A hidden variable $C$ names the component that generated a point, and the density is

For continuous data the natural component is a multivariate Gaussian, giving a mixture of Gaussians with parameters $w_i=P(C=i)$ (the weight), ${\mu }_i$ (the mean), and ${\Sigma }_i$ (the covariance) for each component $i$. If you knew which component generated each point you could fit each Gaussian directly; if you knew the parameters you could assign each point to a component. When neither is known, EM alternates between the two estimates.

Initialize the parameters arbitrarily, then iterate two steps. In the E-step, compute the responsibility $p_{ij}=P(C=i\mid {\mathbf{x}}_j)$, the probability that point ${\mathbf{x}}_j$ came from component $i$, by Bayes' rule from the current parameters. In the M-step, refit each component to all the data, weighting each point by its responsibility.

Writing $n_i={\sum }_j{p}_{ij}$ for the effective number of points assigned to component $i$, the M-step updates are responsibility-weighted versions of the Gaussian ML estimators derived earlier:

The E-step computes the expected values $p_{ij}$ of the hidden indicator variables $Z_{ij}$ ($1$ if point $j$ came from component $i$, else $0$); the M-step maximizes the log likelihood given those expectations. Applied to data sampled from a three-component mixture, EM reconstructs a model nearly indistinguishable from the generator. As pseudocode:

One EM iteration, worked in numbers

Consider a single pass on four one-dimensional points, $x=\{1.0,1.5,5.0,6.0\}$, which plainly form two clusters. Fit a two-component mixture, and initialize deliberately badly to see EM correct itself: means ${\mu }_1=0.0$, ${\mu }_2=4.0$, standard deviations ${\sigma }_1={\sigma }_2=2.0$, and equal weights $w_1=w_2=0.5$.

E-step. For each point compute the responsibility $p_{1j}=P(C=1\mid x_j)$ by Bayes' rule, $p_{1j}=\frac{w_1\mathcal{N}(x_j;{\mu }_1,{\sigma }_1)}{w_1\mathcal{N}(x_j;{\mu }_1,{\sigma }_1)+w_2\mathcal{N}(x_j;{\mu }_2,{\sigma }_2)}$. With $\mathcal{N}(1.0;0,2)=0.1760$ and $\mathcal{N}(1.0;4,2)=0.0648$, the point $x=1.0$ gets $p_{11}=\frac{0.5(0.1760)}{0.5(0.1760)+0.5(0.0648)}=0.731$. The full set of responsibilities to component 1:

The two low points lean toward component 1, the two high points toward component 2, with the boundary points fractionally split — exactly the soft assignment that distinguishes EM from hard $k$-means.

M-step. The effective counts are $n_1={\sum }_j{p}_{1j}=0.731+0.622+0.047+0.018=1.419$ and $n_2=4-n_1=2.581$. The new means are responsibility-weighted averages:

The weighted variances give ${\sigma }_1\leftarrow 0.89$ and ${\sigma }_2\leftarrow 1.88$, and the mixing weights become $w_1\leftarrow n_1/4=0.355$, $w_2\leftarrow 0.645$. In one step the means have jumped from $(0.0,4.0)$ toward the true cluster centers near $1.25$ and $5.5$, and the data log likelihood rose from $-9.40$ to $-7.76$. That increase is guaranteed on every iteration.

Two facts anchor EM's behavior. It increases the log likelihood of the data at every iteration — this can be proved in general — and under mild conditions it converges to a local maximum. It resembles gradient-based hill-climbing but has no step-size parameter to tune.

EM is not foolproof. A component can shrink onto a single point, sending its variance to zero and its likelihood to infinity; two components can merge onto the same data. These degenerate local maxima worsen in high dimensions. Placing priors on the parameters (the MAP version of EM), restarting a collapsing component, and sensible initialization all help.

EM for Bayesian networks and HMMs

The same insight extends beyond mixtures. To learn a Bayesian network with hidden variables, treat the E-step as computing, by ordinary Bayes-net inference, the expected counts that you would otherwise tabulate from complete data. The canonical example is the two-bag candy mixture: two bags of candy have been tipped together, so each piece has a $Flavor$ (cherry or lime), a $Wrapper$ (red or green), and either a $Hole$ or not, but the $Bag$ it came from is hidden. Within a bag the three features are independent — a naive-Bayes model — but the distribution of each feature depends on the bag.

For that model, with the bag hidden, the expected count of candies from bag $1$ is

and each CPT parameter ${\theta }_{ijk}=P(X_i=x_{ij}\mid {\mathbf{U}}_i={\mathbf{u}}_{ik})$ is updated by the ratio of expected counts,

One EM iteration on the candy bags, in numbers

The Gaussian mixture above was carried through with real numbers; the candy net deserves the same. Generate $N=1000$ candies from a true model in which the two bags are equally likely, bag 1 is mostly cherry / red / holed and bag 2 mostly lime / green / no-hole:

Because $Bag$ is hidden, the data are just counts of the eight observable candy types $(Flavor,Wrapper,Holes)$, worked here with the sample counts and initialization AIMA uses:¹

Initialize the parameters deliberately off the truth, ${\theta }^{(0)}=0.6$ and every feature probability $0.6$ for bag 1 and $0.4$ for bag 2.

E-step. For each of the eight candy types, compute the responsibility $P(Bag=1\mid flavor,wrapper,holes)$ by Bayes' rule on the naive-Bayes model,

For the 273 red-wrapped cherry candies with holes, every factor favors bag 1, so

and the eight responsibilities range from $0.835$ down to $0.308$ for the type that best matches bag 2:

M-step. Sum the responsibilities over all candies to get the expected count of bag-1 candies, then divide by $N$ for the new prior. Weighting each type's count by its responsibility (the $273$ candies contribute $273\times 0.835=227.97$, and so on across all eight), the expected bag-1 total is $\hat{N}(Bag=1)=612.4$, so

The feature parameters follow the same ratio-of-expected-counts recipe. For ${\theta }_{F1}=P(cherry\mid Bag=1)$, sum the responsibilities of the four cherry types for the numerator and all eight for the denominator, ${\theta }_{F1}^{(1)}=\hat{N}(cherry,Bag=1)/\hat{N}(Bag=1)$. Doing this for every parameter gives

Every bag-1 parameter has moved up from $0.6$ toward the true $0.8$, and every bag-2 parameter down from $0.4$ toward the true $0.3$ — one iteration already separates the bags in the right direction. The log likelihood of the 1000 candies rises from about $-2044$ under ${\theta }^{(0)}$ to about $-2021$ after this single step, an improvement in the likelihood itself by a factor of roughly $e^{23}\approx 10^{10}$.

Two features of the candy run mirror the general theory. The assignments are soft — the ambiguous types sit at exactly $0.5$ and contribute to both bags — which is why EM, unlike a hard bag-by-bag count, can make progress with the bag entirely unobserved. And each parameter update reads off only a local posterior: ${\theta }_{F1}$ needs the posterior over $Bag$ for each candy and nothing about $Wrapper$ or $Holes$ beyond what that posterior already summarizes. This locality is what turns the update into a by-product of ordinary Bayes-net inference.

The general lesson: the parameter updates for Bayesian-network learning with hidden variables are read directly off the results of inference on each example, and only local posteriors — over each variable and its parents — are needed. Exact inference algorithms such as variable elimination produce these as a by-product, with no learning-specific computation.

The final case is the hidden Markov model. An HMM is a dynamic Bayesian network with one discrete state variable, and each data point is an observation sequence. The complication over general Bayes nets is that the transition probabilities ${\theta }_{ijt}=P(X_{t+1}=j\mid X_t=i)$ are shared across all time steps (${\theta }_{ijt}={\theta }_{ij}$ for all $t$), so the update sums expected transition counts over time,

The expected counts come from the forward–backward algorithm, run as smoothing rather than filtering — you must consider later evidence to estimate that a transition occurred. This instance of EM is the Baum–Welch algorithm.

The general form

Strip the examples away and one equation is left. Let $\mathbf{x}$ be all observed values, $\mathbf{Z}$ all hidden variables, and $\mathbf{\theta }$ all parameters. EM iterates

The E-step is the summation: the expected complete-data log likelihood under the posterior $P(\mathbf{Z}\mid \mathbf{x},{\mathbf{\theta }}^{(i)})$ over the hidden variables. The M-step is the maximization of that expected log likelihood over the parameters. Every instance is this equation with $\mathbf{Z}$ specialized: the component indicator $Z_{ij}$ for mixtures, an unobserved variable's value for Bayes nets, the sequence state for HMMs.

Why EM climbs: the evidence lower bound

The monotone-increase guarantee follows from one inequality. The quantity EM actually wants to raise is the log likelihood of the observed data, $\log P(\mathbf{x}\mid \mathbf{\theta })=\log {\sum }_{\mathbf{z}}P(\mathbf{x},\mathbf{z}\mid \mathbf{\theta })$, a log-of-a-sum that is hard to optimize directly because the hidden $\mathbf{z}$ is buried inside. Introduce any distribution $q(\mathbf{z})$ over the hidden variables and rewrite, then apply Jensen's inequality (the log of an average is at least the average of the logs, since $\log$ is concave):

The right-hand side $\mathcal{L}(q,\mathbf{\theta })$ is the evidence lower bound (ELBO). The gap between it and the true log likelihood coincides with the Kullback–Leibler divergence $\text{KL}(q(\mathbf{z})\Vert P(\mathbf{z}\mid \mathbf{x},\mathbf{\theta }))\ge 0$, so the bound is tight — equality holds — precisely when $q(\mathbf{z})=P(\mathbf{z}\mid \mathbf{x},\mathbf{\theta })$, the posterior over the hidden variables.

EM is coordinate ascent on this bound. The E-step maximizes $\mathcal{L}$ over $q$ with $\mathbf{\theta }$ fixed, and the maximizer is the posterior $q(\mathbf{z})=P(\mathbf{z}\mid \mathbf{x},{\mathbf{\theta }}^{(i)})$ — which is why the E-step computes exactly that posterior, and after it the bound touches the log likelihood. The M-step then maximizes $\mathcal{L}$ over $\mathbf{\theta }$ with $q$ fixed, which (dropping the $\mathbf{\theta }$-free entropy of $q$) is the expected complete-data log likelihood $Q(\mathbf{\theta })$ maximized earlier. Because each step never lowers $\mathcal{L}$ and the E-step makes $\mathcal{L}$ equal the log likelihood, the log likelihood itself cannot decrease:

That chain is the whole convergence proof. It also explains the failure modes: EM maximizes a bound that only touches the true objective at the current parameters, so it can settle at any local maximum, and where the likelihood is unbounded (a Gaussian component collapsing onto one point) EM will climb toward infinity.

Once the general form is understood, variants follow. When the exact E-step is intractable — large Bayes nets — an approximate E-step still yields effective learning: with an MCMC sampler, each configuration of hidden and observed variables visited is treated as a complete observation, and parameters update after each transition. Variational and loopy methods serve the same role for very large networks. The same latent-variable-and-EM story reappears in the deep-learning treatment of latent-variable models, where a neural network parameterizes the components and the E-step becomes an approximate inference network rather than exact Bayes' rule — the variational autoencoder is EM with the posterior itself learned.

From EM to modern probabilistic learning

The candy bags and mixtures of Gaussians in this lesson are the classical core of statistical learning; the same machinery scales into much of modern machine learning, and the public literature fills in the pieces AIMA only sketches.

EM's own history and generality. The algorithm was named and given its general form by Dempster, Laird, and Rubin (1977), who proved the monotone-likelihood result via the lower-bound argument above; the coordinate-ascent-on-the-ELBO view was made explicit by Neal and Hinton (1998), which is what licenses partial or incremental E-steps that update only some responsibilities per pass and still converge.² The Baum–Welch algorithm for HMMs actually predates the general EM formulation (Baum et al., 1970), a special case discovered first.³

Conjugate priors and the exponential family. The beta–binomial and Dirichlet–multinomial conjugacies used here are instances of a general fact: every distribution in the exponential family has a conjugate prior, and Bayesian updating reduces to adding sufficient statistics to the prior's hyperparameters. This structure, laid out in Bishop's Pattern Recognition and Machine Learning (2006), is what makes tractable Bayesian inference possible at all, and it is the backbone of latent Dirichlet allocation (Blei, Ng, and Jordan, 2003), the topic model that applies exactly the mixture-and-EM idea of this lesson to documents, with words drawn from latent topics.⁴

When EM is intractable: variational inference and MCMC. For models where the exact posterior in the E-step cannot be computed, two families of approximation dominate. Variational inference replaces the true posterior with the closest member of a tractable family by maximizing the same ELBO derived above — turning inference into optimization (Jordan et al., 1999; Blei, Kucukelbir, and McAuliffe, 2017).⁵ Markov-chain Monte-Carlo instead draws samples from the posterior; Gelman et al.'s Bayesian Data Analysis is the standard reference. The variational autoencoder (Kingma and Welling, 2013) is the meeting point of both threads and this lesson: it is EM in which a neural network amortizes the E-step, learning a single inference network $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ that approximates the posterior for every data point at once, trained by maximizing the ELBO with the reparameterization trick.⁶ Read that way, the latent-variable models behind modern generative AI are direct descendants of the mixture-of-Gaussians EM worked here.

Learning Bayes-net structure, in practice. The score-and-search sketch of this lesson corresponds to real algorithms: the BIC/MDL score penalizes structure by a term proportional to $\frac{1}{2}(\log N)|\mathbf{\theta }|$, and the constraint-based PC algorithm (Spirtes, Glymour, and Scheines, 1993) recovers structure from conditional-independence tests, a starting point for the modern field of causal discovery.⁷

EM is the standard method for learning when data are incomplete. Whenever a model has variables you cannot observe — clusters without labels, diseases behind symptoms, states behind an emission sequence — the procedure is the same: complete the data in expectation, maximize as if it were real, repeat.