KBKnowledge Base
Machine Learning · 2.2.7

Expectation–Maximization as Optimization

The ELBO, E-step/M-step, and why EM is guaranteed to never get worse.

On this page
In plain English — beginner to advanced

Beginner: imagine data that comes from one of several unknown groups, but you don't know which group each point actually came from — the group membership is a hidden (latent) variable. If you already knew every point's group, fitting each group's parameters would be easy, ordinary maximum likelihood (Module 1, section 2.1.6). If you already knew each group's parameters, guessing which group each point probably came from would also be easy. EM just alternates between these two easy steps until they agree with each other.

Intermediate: name the two steps precisely. The E-step computes, for the current parameter guess, the probability each point belongs to each group — a soft assignment, not a hard decision — by literally applying Bayes' rule (Module 1's Statistical Decision Theory lesson) with the current parameters. The M-step then re-estimates the parameters by ordinary maximum likelihood, but weighting every point's contribution by those soft-assignment probabilities from the E-step.

Advanced: direct MLE is intractable here because the likelihood you actually want to maximize has a SUM (over all possible hidden-variable values) sitting inside a logarithm, log Σ_z P(x,z|θ) — contrast this with the sum-of-logs, Σ log P(xᵢ|θ), from an ordinary MLE problem, which was easy precisely because differentiating a sum of logs is trivial while differentiating a log of a sum has no clean closed form in general. EM is a way to make guaranteed, monotonic progress on this hard objective without ever having to differentiate it directly.

Formula
logP(xθ)=ELBO(q,θ)+DKL(q(z)P(zx,θ))\log P(x|\theta) = \text{ELBO}(q,\theta) + D_{KL}\big(q(z)\,\|\,P(z|x,\theta)\big)
ELBO(q,θ)=Eq(z)[logP(x,zθ)]Eq(z)[logq(z)]\text{ELBO}(q,\theta) = \mathbb{E}_{q(z)}[\log P(x,z|\theta)] - \mathbb{E}_{q(z)}[\log q(z)]

The KL term is always ≥ 0, so the ELBO is always a valid lower bound on the true log-likelihood — hence "Evidence Lower BOund."

Derivation: the ELBO decomposition, and why EM never makes the likelihood worse

Start from the true log-likelihood and multiply/divide inside by an arbitrary distribution q(z) over the hidden variable:

logP(xθ)=logzq(z)P(x,zθ)q(z)\log P(x|\theta) = \log \sum_z q(z)\frac{P(x,z|\theta)}{q(z)}

Apply Jensen's inequality — for a concave function like log, log E[Y] ≥ E[log Y] — to pull the log inside the expectation over q:

logP(xθ)zq(z)logP(x,zθ)q(z)=ELBO(q,θ)\log P(x|\theta) \ge \sum_z q(z)\log\frac{P(x,z|\theta)}{q(z)} = \text{ELBO}(q,\theta)

That's the lower-bound half. Now derive the EXACT gap. Substitute the definitions of ELBO and KL divergence and simplify directly:

logP(xθ)ELBO(q,θ)=zq(z)logq(z)P(zx,θ)=DKL(q(z)P(zx,θ))\log P(x|\theta) - \text{ELBO}(q,\theta) = \sum_z q(z)\log\frac{q(z)}{P(z|x,\theta)} = D_{KL}\big(q(z)\,\|\,P(z|x,\theta)\big)

The true log-likelihood minus the ELBO is EXACTLY the KL divergence between whatever q you chose and the TRUE posterior P(z|x,θ). Since KL divergence is minimized — equal to exactly zero — precisely when q equals that true posterior, choosing q(z) = P(z|x,θ_old) in the E-step makes the ELBO exactly EQUAL to the true log-likelihood at the current θ_old: the bound "touches" the true curve exactly at the current point.

Now the M-step: with q fixed from the E-step, maximize the ELBO over θ alone. Chain the inequalities:

  1. The ELBO touches the true log-likelihood at θ_old: ELBO(q,θ_old) = log P(x|θ_old).
  2. The M-step picks θ_new to maximize the ELBO, so ELBO(q,θ_new) ≥ ELBO(q,θ_old).
  3. The ELBO is a lower bound EVERYWHERE, including at θ_new: log P(x|θ_new) ≥ ELBO(q,θ_new).

Chaining all three: log P(x|θ_new) ≥ ELBO(q,θ_new) ≥ ELBO(q,θ_old) = log P(x|θ_old). The true log-likelihood at the new parameters is at least as large as at the old ones — EM can never make the likelihood worse, on every single iteration, guaranteed.

Where this is used: this exact E-step/M-step alternation, applied to a mixture-of-Gaussians likelihood, is Gaussian Mixture Model fitting — covered directly in a later Clustering module of this chapter, using precisely the machinery derived here.

EM climbing the true log-likelihood, one touching lower bound at a time

Each dashed violet curve touches the true (solid black) curve at the current θ and lies below it everywhere else. Jumping to that curve's own maximum and drawing a new one there can only move the true log-likelihood up or hold it flat — never down. Drag the red dot to restart from a different θ.

Practical example — EM for a 1D two-component Gaussian mixture
python
import numpy as np

def gaussian_pdf(x, mu, var):
    return np.exp(-((x - mu) ** 2) / (2 * var)) / np.sqrt(2 * np.pi * var)

def em_gmm_1d(x, iters=50):
    n = len(x)
    mu = np.array([x.min(), x.max()])
    var = np.array([x.var(), x.var()])
    weight = np.array([0.5, 0.5])
    loglik_history = []

    for _ in range(iters):
        # E-step: responsibilities via Bayes' rule with the current parameters.
        p0 = weight[0] * gaussian_pdf(x, mu[0], var[0])
        p1 = weight[1] * gaussian_pdf(x, mu[1], var[1])
        total = p0 + p1
        r0, r1 = p0 / total, p1 / total

        # M-step: responsibility-weighted re-estimation.
        n0, n1 = r0.sum(), r1.sum()
        mu[0], mu[1] = (r0 * x).sum() / n0, (r1 * x).sum() / n1
        var[0] = (r0 * (x - mu[0]) ** 2).sum() / n0
        var[1] = (r1 * (x - mu[1]) ** 2).sum() / n1
        weight[0], weight[1] = n0 / n, n1 / n

        loglik = np.log(total).sum()
        loglik_history.append(loglik)

    return mu, var, weight, loglik_history

rng = np.random.default_rng(0)
x = np.concatenate([rng.normal(2, 1, 150), rng.normal(8, 1.5, 150)])
mu, var, weight, history = em_gmm_1d(x)
print("fitted means:", np.round(mu, 2), " (true: ~2, ~8)")
print("log-likelihood is non-decreasing:", all(b >= a - 1e-9 for a, b in zip(history, history[1:])))
Real-world examples
  • Gaussian Mixture Models for soft clustering — the direct next use of this exact machinery, in a later Clustering module of this chapter.
  • Hidden Markov Model parameter estimation via Baum-Welch is itself an application of EM, covered in a later Graphical Models module.
  • Missing-data imputation problems are often framed as exactly a latent-variable MLE problem and solved via EM.
  • Topic models (Latent Dirichlet Allocation, a later module) use EM-like alternating inference between topic assignments and topic-word distributions.
  • Item-response theory / psychometric models in educational testing use EM to jointly estimate latent student ability and item difficulty — a case where the "hidden variable" is a genuinely unobservable real-world quantity, not just a clustering convenience.
Common mistakes
  • Assuming EM's monotonic increase in likelihood means it finds the GLOBAL maximum — it only guarantees non-decreasing progress toward SOME local maximum. Different initializations can converge to different, sometimes much worse, local optima, because the marginal likelihood with a latent variable is generically not convex in θ even when the complete-data likelihood would have been easy (Module 2's Convex Optimization Basics lesson).
  • Stopping EM based on parameter change alone rather than log-likelihood (or ELBO) change, which can be a less reliable convergence signal.
  • Expecting fast convergence near a plateau — EM can crawl very slowly there even while technically still monotonically improving, which in practice sometimes calls for accelerated or hybrid variants.
Going deeper

Variational inference (a later Probabilistic & Bayesian Methods module) is a direct generalization of this exact ELBO idea: instead of restricting q to be the exact posterior (which may be intractable for more complex models), it optimizes the ELBO over some restricted, tractable FAMILY of q distributions (mean-field, etc.). EM is the special case where that family is rich enough to contain the true posterior exactly, so the E-step can always close the gap completely; when it can't, the result is "variational EM" — a strictly more general algorithm built on the identical machinery derived in this lesson.

Check yourself
Why does EM's guaranteed monotonic increase in the true log-likelihood NOT imply it reaches the global maximum?

The proof only chains three inequalities that establish log P(x|θ_new) >= log P(x|θ_old) — it never claims θ_new is anywhere near the best possible θ overall, only that it's no worse than where you started. Since the marginal likelihood with a latent variable is generally non-convex, there can be multiple local maxima, and which one EM climbs to depends entirely on where it started. This is why GMM fitting in practice is typically run from several random initializations and the best resulting likelihood is kept.

Key takeaway

With convexity (2.2.1), gradient-based methods (2.2.2-2.2.3), second-order methods (2.2.4), proximal methods (2.2.5), duality (2.2.6), and now EM (2.2.7) in hand, this module's toolbox is complete: every specific model in the rest of this Machine Learning chapter is simply a choice of hypothesis class and loss function (Module 1) paired with whichever of these optimization tools fits that objective's particular shape.

Newsletter

Stay in the loop

Subscribe to get new docs, diagrams, and engineering write-ups by Dharaneesh Boobalan delivered to your inbox.

  • Deep-dive write-ups on ML, inference, and systems.
  • New Draw.io diagrams & interactive canvases.
  • Agentic patterns and rocket-science notes.
  • No spam. One tasteful email when there's something new.

Crafted by Dharaneesh Boobalan

Newsletter

Get new docs, diagrams, and write-ups in your inbox.

We never share your details. Unsubscribe anytime.