The EM Algorithm

Why This Matters

The EM algorithm is one of the most important algorithms in statistical machine learning. Whenever your model contains latent (unobserved) variables, you cannot simply write down the log-likelihood and differentiate. the latent variables couple the parameters in a way that makes direct optimization intractable. EM sidesteps this by alternating between inferring the latent variables and updating the parameters.

Gaussian mixture models, hidden Markov models, factor analysis, probabilistic PCA, topic models. all of these are classically fitted with EM. And the core idea. maximizing a lower bound on the log-likelihood. is exactly the principle behind variational autoencoders and modern variational inference.

If you understand EM deeply, you understand the conceptual engine behind a huge fraction of probabilistic machine learning.

Mental Model

You want to maximize the log-likelihood $\log p(x \mid \theta)$, but there are hidden variables $z$ that make this hard. Direct marginalization $\log \sum_z p(x, z \mid \theta)$ puts a sum inside a log, which is analytically and computationally painful.

EM's trick: instead of maximizing the log-likelihood directly, construct a lower bound (the ELBO) that is easy to compute and maximize. Iterate: tighten the bound (E-step), then maximize it (M-step). Each iteration is guaranteed to increase (or not decrease) the true log-likelihood.

The Missing-Data Formulation

Definition

Complete-Data Log-Likelihood

The complete-data log-likelihood is $\log p(x, z \mid \theta)$, where $x$ is the observed data and $z$ is the latent (missing) data. If we knew $z$, optimization over $\theta$ would typically be tractable: for exponential-family models it reduces to sufficient-statistic matching with a closed-form solution.

Definition

Incomplete-Data Log-Likelihood

The incomplete-data log-likelihood (or marginal log-likelihood) is:

$\log p(x \mid \theta) = \log \sum_z p(x, z \mid \theta)$

This is what we actually want to maximize, but the sum (or integral) inside the log makes it hard.

The fundamental problem: we want to maximize $\log p(x \mid \theta)$, but we can only easily work with $\log p(x, z \mid \theta)$. EM bridges this gap.

The ELBO Derivation

This is the mathematical heart of EM. For any distribution $q(z)$ over the latent variables:

$\log p(x \mid \theta) = \log \sum_z p(x, z \mid \theta)$

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

$\geq \sum_z q(z) \log \frac{p(x, z \mid \theta)}{q(z)}$

where the inequality is Jensen's inequality applied to the concave function $\log(\cdot)$.

Definition

Evidence Lower Bound (ELBO)

The ELBO (Evidence Lower BOund) is:

$\mathcal{L}(q, \theta) = \sum_z q(z) \log \frac{p(x, z \mid \theta)}{q(z)} = \mathbb{E}_{q}[\log p(x, z \mid \theta)] - \mathbb{E}_{q}[\log q(z)]$

It is a lower bound on $\log p(x \mid \theta)$ for any $q$. Equivalently:

$\mathcal{L}(q, \theta) = \mathbb{E}_{q}[\log p(x, z \mid \theta)] + H(q)$

where $H(q) = -\sum_z q(z) \log q(z)$ is the entropy of $q$.

Theorem

ELBO Decomposition

Statement

For any distribution $q(z)$ and parameters $\theta$:

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

where $D_{\mathrm{KL}}$ is the Kullback-Leibler divergence.

Intuition

The log-likelihood decomposes exactly into the ELBO plus the KL divergence between $q(z)$ and the true posterior $p(z \mid x, \theta)$. Since KL divergence is always $\geq 0$, the ELBO is always $\leq \log p(x \mid \theta)$. The bound is tight (equality holds) when $q(z) = p(z \mid x, \theta)$. that is, when $q$ equals the true posterior.

Proof Sketch

Start from the ELBO definition and expand:

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

Rearranging gives the decomposition.

Why It Matters

This decomposition is the raison d'etre of EM. It tells you that to make the ELBO as tight as possible (i.e., equal to the log-likelihood), you should set $q(z) = p(z \mid x, \theta)$. That is exactly what the E-step does. Then maximizing the ELBO over $\theta$ (the M-step) increases the log-likelihood.

Failure Mode

If the true posterior $p(z \mid x, \theta)$ is intractable to compute, you cannot run the E-step exactly. This is the setting where you need variational EM (restrict $q$ to a tractable family) or Monte Carlo EM (approximate the expectation by sampling).

The EM Algorithm

With the ELBO in hand, EM is simply coordinate ascent on $\mathcal{L}(q, \theta)$:

Initialize $\theta^{(0)}$.

E-step (iteration $t$): Set $q^{(t)}(z) = p(z \mid x, \theta^{(t)})$.

This makes the bound tight: $\mathcal{L}(q^{(t)}, \theta^{(t)}) = \log p(x \mid \theta^{(t)})$.

M-step (iteration $t$): Set $\theta^{(t+1)} = \arg\max_{\theta} \mathcal{L}(q^{(t)}, \theta)$.

Since $q^{(t)}$ is now fixed, this is equivalent to:

$\theta^{(t+1)} = \arg\max_{\theta} \mathbb{E}_{q^{(t)}}[\log p(x, z \mid \theta)]$

because $H(q^{(t)})$ does not depend on $\theta$.

The quantity $Q(\theta, \theta^{(t)}) = \mathbb{E}_{p(z \mid x, \theta^{(t)})}[\log p(x, z \mid \theta)]$ is called the Q-function in the classical EM literature.

Monotonic Convergence

Theorem

EM Monotonically Increases the Likelihood

Statement

At each iteration of EM:

$\log p(x \mid \theta^{(t+1)}) \geq \log p(x \mid \theta^{(t)})$

with equality if and only if $\theta^{(t)}$ is already a fixed point of EM.

Intuition

The E-step makes the ELBO equal to the log-likelihood. The M-step increases the ELBO (or keeps it the same). Since the ELBO was equal to the log-likelihood at $\theta^{(t)}$, and we increased it, the new log-likelihood at $\theta^{(t+1)}$ must be at least as large.

Proof Sketch

After the E-step: $\mathcal{L}(q^{(t)}, \theta^{(t)}) = \log p(x \mid \theta^{(t)})$.

After the M-step: $\mathcal{L}(q^{(t)}, \theta^{(t+1)}) \geq \mathcal{L}(q^{(t)}, \theta^{(t)})$ by definition of maximization.

But $\mathcal{L}(q^{(t)}, \theta^{(t+1)}) \leq \log p(x \mid \theta^{(t+1)})$ because the ELBO is always a lower bound.

Chaining: $\log p(x \mid \theta^{(t+1)}) \geq \mathcal{L}(q^{(t)}, \theta^{(t+1)}) \geq \mathcal{L}(q^{(t)}, \theta^{(t)}) = \log p(x \mid \theta^{(t)})$

Why It Matters

Monotonic convergence is a strong guarantee: EM never makes things worse. Combined with the fact that the log-likelihood is bounded above (for well-defined models), this ensures that EM converges to some stationary point. But it does not guarantee convergence to the global maximum.

Failure Mode

If the M-step only partially maximizes the Q-function (generalized EM), you still get monotonic convergence as long as $Q(\theta^{(t+1)}, \theta^{(t)}) \geq Q(\theta^{(t)}, \theta^{(t)})$. But convergence can be very slow, especially near saddle points.

Why Jensen's Inequality Is the Key

Jensen's inequality states that for a concave function $f$:

$f(\mathbb{E}[X]) \geq \mathbb{E}[f(X)]$

In the ELBO derivation, we apply this to $f = \log$, which is concave. The "random variable" is $p(x, z \mid \theta) / q(z)$, and the expectation is with respect to $q(z)$.

Without Jensen's inequality, we would have no lower bound, no ELBO, and no EM. The entire algorithm rests on the concavity of the logarithm.

Canonical Example: Gaussian Mixture Models

Example

EM for Gaussian Mixture Models

A GMM models data as arising from $K$ Gaussian components:

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

where $\theta = \{\pi_k, \mu_k, \Sigma_k\}_{k=1}^K$ and $\pi_k$ are mixing weights with $\sum_k \pi_k = 1$.

The latent variable $z_i \in \{1, \ldots, K\}$ indicates which component generated datapoint $x_i$.

E-step: Compute responsibilities (posterior over component assignments):

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

M-step: Update parameters using the responsibilities as soft assignments:

$N_k = \sum_{i=1}^n \gamma_{ik}$

$\mu_k^{(t+1)} = \frac{1}{N_k} \sum_{i=1}^n \gamma_{ik} \, x_i$

$\Sigma_k^{(t+1)} = \frac{1}{N_k} \sum_{i=1}^n \gamma_{ik} (x_i - \mu_k^{(t+1)})(x_i - \mu_k^{(t+1)})^\top$

$\pi_k^{(t+1)} = \frac{N_k}{n}$

Each iteration is guaranteed to increase (or maintain) the log-likelihood. The algorithm converges to a local maximum of the GMM likelihood.

HMMs and the Baum-Welch Algorithm

The Baum-Welch algorithm for training hidden Markov models is a special case of EM. The E-step uses the forward-backward algorithm to compute posterior marginals $p(z_t \mid x_{1:T}, \theta)$ and pairwise posteriors $p(z_t, z_{t+1} \mid x_{1:T}, \theta)$ over the hidden states. The M-step re-estimates transition probabilities and emission parameters using these posteriors as soft counts.

The structure of the HMM makes the E-step tractable despite an exponential number of latent sequences, because the forward-backward algorithm exploits the chain structure via dynamic programming.

Common Confusions

Watch Out

EM does NOT find global optima

EM finds local optima (more precisely, stationary points) of the log-likelihood. The monotonic convergence theorem guarantees that the likelihood never decreases, but it says nothing about reaching the global maximum. For multi-modal likelihoods (which GMMs almost always have), the solution depends heavily on initialization. In practice, people run EM with multiple random restarts and take the solution with highest likelihood.

Watch Out

The E-step is expectation, not sampling

A common misconception is that the E-step involves sampling from the posterior $p(z \mid x, \theta)$. It does not. The E-step computes $\mathbb{E}_{p(z \mid x, \theta^{(t)})}[\log p(x, z \mid \theta)]$ as a function of $\theta$. For discrete latent variables, this means computing exact posterior probabilities. For continuous latent variables, it means computing posterior moments analytically. If you sample instead of computing exact expectations, you are doing Monte Carlo EM, which is a different (stochastic) algorithm with different convergence properties.

Watch Out

EM is not just for mixtures

While GMMs are the canonical example, EM applies to any model with latent variables: factor analysis, probabilistic PCA, hidden Markov models, latent Dirichlet allocation (with variational E-step), missing data imputation, and many more. The missing-data formulation is the general framework.

Summary

• EM maximizes a lower bound (the ELBO) on the log-likelihood
• The E-step sets $q(z) = p(z \mid x, \theta^{(t)})$, making the bound tight
• The M-step maximizes $\mathbb{E}_q[\log p(x, z \mid \theta)]$ over $\theta$
• Each iteration monotonically increases the log-likelihood
• Convergence is to a local optimum, not necessarily the global optimum
• Jensen's inequality on the concave $\log$ function is the foundational step
• For exponential family models, the M-step often has closed-form solutions

Exercises

ExerciseCore

Problem

Consider a two-component GMM in 1D with equal mixing weights $\pi_1 = \pi_2 = 0.5$, known variance $\sigma^2 = 1$, and unknown means $\mu_1, \mu_2$. You observe a single datapoint $x = 0$.

(a) Write the E-step: compute $\gamma_1 = p(z = 1 \mid x = 0, \mu_1^{(t)}, \mu_2^{(t)})$.

(b) Write the M-step update for $\mu_1$ and $\mu_2$ given $n$ datapoints.

(c) If $\mu_1^{(0)} = -1$ and $\mu_2^{(0)} = 1$, what are $\gamma_1$ and $\gamma_2$ for $x = 0$?

Hint

Reveal Solution

ExerciseAdvanced

Problem

Prove that if the M-step only partially maximizes the Q-function. that is, $Q(\theta^{(t+1)}, \theta^{(t)}) \geq Q(\theta^{(t)}, \theta^{(t)})$ but $\theta^{(t+1)}$ is not necessarily the global maximizer. The monotonic convergence property still holds. This is the Generalized EM (GEM) algorithm.

Hint

Reveal Solution

References

Canonical:

• Dempster, Laird, Rubin, "Maximum Likelihood from Incomplete Data via the EM Algorithm" (1977). The original paper
• Bishop, Pattern Recognition and Machine Learning, Chapter 9
• McLachlan & Krishnan, The EM Algorithm and Extensions (2008)

Current:

• Murphy, Probabilistic Machine Learning: An Introduction (2022), Chapter 8

• Neal & Hinton, "A View of the EM Algorithm that Justifies Incremental, Sparse, and Other Variants" (1998)

• Hastie, Tibshirani, Friedman, The Elements of Statistical Learning (2009)

Next Topics

The natural next steps from EM:

• EM algorithm variants: Monte Carlo EM, variational EM, stochastic EM, and generalized EM for intractable E-steps and M-steps