The EM Algorithm

Setting Up the Problem

Suppose we have observed data $\mathbf{X} = \{x_1, \ldots, x_N\}$ and postulate that each $x_i$ was generated alongside an unobserved latent variable $z_i$ drawn from some distribution. The complete-data log-likelihood would be easy to optimize if we knew $\mathbf{Z} = \{z_1, \ldots, z_N\}$, but we do not. The observed-data log-likelihood is:

Explanation: The log-likelihood of the observed data requires summing (or integrating) over all possible configurations of the latent variables, which moves the sum inside the log and makes direct maximization analytically painful. Symbols: $\theta$ is the vector of all model parameters; $p(\mathbf{X}, \mathbf{Z} \mid \theta)$ is the joint density of observed and latent data; the sum becomes an integral for continuous latents.

The EM Lower Bound (ELBO)

Introduce any distribution $q(\mathbf{Z})$ over latent variables. By Jensen's inequality applied to the concave log function:

Explanation: This is the Evidence Lower BOund (ELBO); the gap between the true log-likelihood and this bound is exactly the KL divergence from $q(\mathbf{Z})$ to the true posterior $p(\mathbf{Z} \mid \mathbf{X}, \theta)$. Symbols: $\mathcal{L}(q, \theta)$ is the lower bound; $q(\mathbf{Z})$ is a free variational distribution we choose; the ratio inside the log separates into the complete-data log-likelihood minus the entropy of $q$.

E-Step: Tighten the Bound

Set $q(\mathbf{Z}) = p(\mathbf{Z} \mid \mathbf{X}, \theta^{\text{old}})$, the posterior over latents given current parameters. This makes the KL gap zero, so the bound touches the true likelihood. Define the Q-function:

Explanation: The Q-function is simply the expected complete-data log-likelihood, where the expectation is taken with respect to the posterior distribution of the latent variables computed under the old parameters. Symbols: $\mathbb{E}_{\mathbf{Z} \mid \mathbf{X}, \theta^{\text{old}}}[\cdot]$ denotes expectation over the posterior of $\mathbf{Z}$; $\theta^{\text{old}}$ is held fixed during this step; $\theta$ is the free variable we will optimize in M-step.

M-Step: Maximize the Bound

Explanation: With the posterior over latents fixed from the E-step, we maximize the expected complete-data log-likelihood over the model parameters, which is typically an analytically tractable problem.

Convergence Guarantee

The monotonic non-decrease property follows directly from the bound structure:

Explanation: Each EM iteration increases (or leaves unchanged) the true observed-data log-likelihood because the M-step increases Q, and Q lower-bounds ℓ tightly at the current parameter value. ---

Worked Example: Gaussian Mixture Models in PyTorch

A Gaussian Mixture Model (GMM) with $K$ components is one of the canonical EM applications. Each data point $x_i \in \mathbb{R}^D$ is assumed drawn from one of $K$ Gaussians with mixing weights $\pi_k$, means $\mu_k$, and covariances $\Sigma_k$. ```python import torch import torch.nn.functional as F class GMMwithEM: def __init__(self, K: int, D: int, n_iter: int = 100): self.K, self.D, self.n_iter = K, D, n_iter

# Random initialisation — run multiple times! self.pi = torch.ones(K) / K # mixing weights self.mu = torch.randn(K, D) # means self.Sigma = torch.stack([torch.eye(D)] * K) # covariances def _log_gaussian(self, X: torch.Tensor) -> torch.Tensor: """Log N(x | mu_k, Sigma_k) for every (n, k) pair. Returns [N, K].""" N = X.shape[0] log_probs = torch.zeros(N, self.K) for k in range(self.K): dist = torch.distributions.MultivariateNormal( self.mu[k], covariance_matrix=self.Sigma[k] ) log_probs[:, k] = dist.log_prob(X) return log_probs # [N, K] def e_step(self, X: torch.Tensor) -> torch.Tensor: """Compute responsibilities r[n, k] = p(z=k | x_n, theta_old).""" log_probs = self._log_gaussian(X) # [N, K] log_weighted = log_probs + torch.log(self.pi).unsqueeze(0) # [N, K]

# Numerically stable softmax over K log_r = log_weighted - torch.logsumexp(log_weighted, dim=1, keepdim=True) return log_r.exp() # responsibilities [N, K] def m_step(self, X: torch.Tensor, r: torch.Tensor) -> None: """Update pi, mu, Sigma from responsibilities.""" N_k = r.sum(dim=0) # effective counts [K] self.pi = N_k / N_k.sum() self.mu = (r.T @ X) / N_k.unsqueeze(1) # [K, D] for k in range(self.K): diff = X - self.mu[k] # [N, D] self.Sigma[k] = (r[:, k].unsqueeze(1) diff).T @ diff / N_k[k] self.Sigma[k] += 1e-6 torch.eye(self.D) # numerical stability def fit(self, X: torch.Tensor) -> list[float]: log_likelihoods = [] for _ in range(self.n_iter): r = self.e_step(X) self.m_step(X, r)

# Track observed log-likelihood log_probs = self._log_gaussian(X) + torch.log(self.pi) ll = torch.logsumexp(log_probs, dim=1).sum().item() log_likelihoods.append(ll) return log_likelihoods

# Usage torch.manual_seed(0) X = torch.cat([torch.randn(200, 2) + t for t in [2., -2.]]) # two clusters model = GMMwithEM(K=2, D=2, n_iter=50) lls = model.fit(X) print(f"Final log-likelihood: {lls[-1]:.2f}")

# ll should increase monotonically — assert this during debugging! assert all(lls[i+1] >= lls[i] - 1e-4 for i in range(len(lls)-1)), "EM violated monotonicity!" ```

Maximum Likelihood Estimation · Gaussian Mixture Models · Variational Inference · Hidden Markov Models · Jensen's Inequality and Convexity