🧮 Brain Teaser

The EM Algorithm's Hidden Monotonicity

You are fitting a latent variable model with observed data XX, latent variables ZZ, and parameters θ\theta. The EM algorithm alternates between:

  • E-step: Compute Q(θ,θold)=EZX,θold[logp(X,Zθ)]Q(\theta, \theta^{\text{old}}) = \mathbb{E}_{Z|X,\theta^{\text{old}}}[\log p(X, Z \mid \theta)]
  • M-step: Set θnew=argmaxθQ(θ,θold)\theta^{\text{new}} = \arg\max_\theta Q(\theta, \theta^{\text{old}})

Prove that the observed-data log-likelihood (θ)=logp(Xθ)\ell(\theta) = \log p(X \mid \theta) is non-decreasing at every EM step:

(θnew)(θold).\ell(\theta^{\text{new}}) \geq \ell(\theta^{\text{old}}).

The key insight is not about optimization tricks — it is a direct consequence of a classical inequality. What is that inequality, and why does it apply here?

EM algorithmJensen's inequalityKL divergencelatent variableslikelihood

Answer: EM Algorithm's Hidden Monotonicity

Key Idea / Intuition

The observed-data likelihood p(Xθ)p(X \mid \theta) is obtained by marginalizing out the latent variables ZZ. The EM algorithm never directly maximizes (θ)\ell(\theta) — instead it maximizes a lower bound on (θ)\ell(\theta) that is tight at θold\theta^{\text{old}}. This lower bound is constructed using Jensen's inequality applied to the concave log\log function. Because the M-step raises (or maintains) the lower bound, and the bound was tight at θold\theta^{\text{old}}, the actual likelihood must go up.


Formal Proof / Solution

Step 1: Decompose the observed log-likelihood

By Bayes' theorem, for any distribution q(Z)q(Z) over the latent variables:

logp(Xθ)=logp(X,Zθ)p(ZX,θ).\log p(X \mid \theta) = \log \frac{p(X, Z \mid \theta)}{p(Z \mid X, \theta)}.

Take expectation under q(Z)=p(ZX,θold)q(Z) = p(Z \mid X, \theta^{\text{old}}) on both sides. The left side doesn't depend on ZZ, so:

(θ)=EZX,θold[logp(X,Zθ)p(ZX,θ)].\ell(\theta) = \mathbb{E}_{Z \mid X, \theta^{\text{old}}}\left[\log \frac{p(X, Z \mid \theta)}{p(Z \mid X, \theta)}\right].

Write this as:

(θ)=EZX,θold[logp(X,Zθ)]Q(θ,θold)EZX,θold[logp(ZX,θ)]H(θ,θold).\ell(\theta) = \underbrace{\mathbb{E}_{Z \mid X, \theta^{\text{old}}}[\log p(X, Z \mid \theta)]}_{Q(\theta,\, \theta^{\text{old}})} - \underbrace{\mathbb{E}_{Z \mid X, \theta^{\text{old}}}[\log p(Z \mid X, \theta)]}_{H(\theta,\, \theta^{\text{old}})}.

Step 2: Bound the HH term using Jensen

Consider the change in HH when moving from θold\theta^{\text{old}} to θnew\theta^{\text{new}}:

H(θnew,θold)H(θold,θold)=EZX,θold[logp(ZX,θnew)p(ZX,θold)].H(\theta^{\text{new}}, \theta^{\text{old}}) - H(\theta^{\text{old}}, \theta^{\text{old}}) = \mathbb{E}_{Z \mid X, \theta^{\text{old}}}\left[\log \frac{p(Z \mid X, \theta^{\text{new}})}{p(Z \mid X, \theta^{\text{old}})}\right].

By Jensen's inequality applied to the concave log\log:

E[logp(ZX,θnew)p(ZX,θold)]logE[p(ZX,θnew)p(ZX,θold)]=logp(ZX,θnew)p(ZX,θold)p(ZX,θold)dZ=log1=0.\mathbb{E}\left[\log \frac{p(Z \mid X, \theta^{\text{new}})}{p(Z \mid X, \theta^{\text{old}})}\right] \leq \log \mathbb{E}\left[\frac{p(Z \mid X, \theta^{\text{new}})}{p(Z \mid X, \theta^{\text{old}})}\right] = \log \int \frac{p(Z \mid X, \theta^{\text{new}})}{p(Z \mid X, \theta^{\text{old}})} p(Z \mid X, \theta^{\text{old}})\, dZ = \log 1 = 0.

This is precisely the statement that KL divergence is non-negative:

KL ⁣(p(ZX,θold)    p(ZX,θnew))0.\text{KL}\!\left(p(Z \mid X, \theta^{\text{old}}) \;\|\; p(Z \mid X, \theta^{\text{new}})\right) \geq 0.

So: H(θnew,θold)H(θold,θold)H(\theta^{\text{new}}, \theta^{\text{old}}) \leq H(\theta^{\text{old}}, \theta^{\text{old}}).

Step 3: Combine

(θnew)(θold)=[Q(θnew,θold)Q(θold,θold)]0 (M-step maximizes Q)[H(θnew,θold)H(θold,θold)]0 (KL ≥ 0)0.\ell(\theta^{\text{new}}) - \ell(\theta^{\text{old}}) = \underbrace{\bigl[Q(\theta^{\text{new}}, \theta^{\text{old}}) - Q(\theta^{\text{old}}, \theta^{\text{old}})\bigr]}_{\geq\, 0 \text{ (M-step maximizes } Q)} - \underbrace{\bigl[H(\theta^{\text{new}}, \theta^{\text{old}}) - H(\theta^{\text{old}}, \theta^{\text{old}})\bigr]}_{\leq\, 0 \text{ (KL ≥ 0)}} \geq 0.

Both terms work in our favour: the M-step ensures QQ goes up, and Jensen/KL ensures HH goes down (or stays).

Summary of the beautiful structure

| Term | Direction | Reason | |------|-----------|--------| | Q(θnew)Q(θold)Q(\theta^{\text{new}}) - Q(\theta^{\text{old}}) | 0\geq 0 | M-step definition | | H(θnew)H(θold)H(\theta^{\text{new}}) - H(\theta^{\text{old}}) | 0\leq 0 | Jensen / KL divergence 0\geq 0 |

Hence (θnew)(θold)\ell(\theta^{\text{new}}) \geq \ell(\theta^{\text{old}}). \blacksquare

The deeper message: EM works by maximizing a variational lower bound (often called the ELBO). The KL non-negativity is precisely what keeps the bound valid, and Jensen is the engine behind KL non-negativity.

Source: The Elements of Statistical Learning, Hastie, Tibshirani et al., Chapter 8

Type: ML/StatsSource: The Elements of Statistical Learning, Hastie, Tibshirani et al., Chapter 8Edit on GitHub ↗