Expectation-Maximization, from Latent Variables to Speech
In this post, I want to give an overview of the Expectation-Maximization (EM) algorithm through two classical examples. EM appears in many older statistical models, but it also sits behind a surprising number of modern latent-variable methods, so it is worth having a clean mental picture of what it is doing.
1. The core algorithm
Suppose we observe data X and we want to estimate parameters \theta by maximum likelihood: \hat{\theta} = \arg\max_\theta p(X \mid \theta).
Maximum likelihood is attractive for several reasons:
- It is consistent: as the amount of data grows, the estimator converges to the true parameter, under suitable regularity assumptions.
- It is asymptotically normal: for large samples, the fluctuations of the estimator around the true parameter are approximately Gaussian.
- It is asymptotically efficient: under suitable regularity assumptions, its asymptotic covariance is the inverse Fisher information.
So maximum likelihood is not just an intuitive criterion. It also comes with a strong theoretical justification. The difficulty is that in latent-variable models the marginal likelihood p(X \mid \theta) = \int p(X, z \mid \theta)\,dz is often hard to optimize directly.
This is the setting where EM is useful. We introduce latent variables Z and build an iterative scheme that is easier to optimize at each step.
Starting from a current iterate \theta^{(n)}, the goal is to find a new parameter value \theta such that \log p(X \mid \theta) \ge \log p(X \mid \theta^{(n)}). In other words, we would like to increase the observed-data log-likelihood at every iteration.
To see how EM does this, let us insert the latent variables into the problem. We write \begin{align*} \log p(X \mid \theta) - \log p(X \mid \theta^{(n)}) &= \log \int p(X, z \mid \theta)\,dz - \log p(X \mid \theta^{(n)}) \\ &= \log \int p(z \mid X, \theta^{(n)}) \frac{p(X, z \mid \theta)}{p(z \mid X, \theta^{(n)})}\,dz - \log p(X \mid \theta^{(n)}) \\ &\ge \int p(z \mid X, \theta^{(n)}) \log \frac{p(X, z \mid \theta)}{p(z \mid X, \theta^{(n)})}\,dz - \log p(X \mid \theta^{(n)}) \end{align*} by Jensen’s inequality.
This is the key step. We have taken the difficult quantity \log p(X \mid \theta) and lower bounded it by an expectation under the posterior distribution of the latent variables at the current iterate \theta^{(n)}.
If we expand the joint density as p(X, z \mid \theta) = p(X \mid z, \theta)p(z \mid \theta), the bound becomes \int p(z \mid X, \theta^{(n)}) \log \frac{p(X \mid z, \theta)p(z \mid \theta)} {p(z \mid X, \theta^{(n)})p(X \mid \theta^{(n)})}\,dz.
Now the terms p(z \mid X, \theta^{(n)}) and p(X \mid \theta^{(n)}) do not depend on \theta. So, for the purpose of maximizing over \theta, they can be treated as constants. This is why maximizing the lower bound is equivalent to maximizing Q(\theta \mid \theta^{(n)}) \coloneqq \mathbb E_{p(\cdot \mid X, \theta^{(n)})}\big[\log p(X, Z \mid \theta)\big].
A useful sanity check is the following. If we define the lower bound on the log-likelihood difference by \widetilde Q(\theta \mid \theta^{(n)}) \coloneqq \mathbb E_{p(\cdot \mid X, \theta^{(n)})} \left[ \log \frac{p(X, Z \mid \theta)} {p(Z \mid X, \theta^{(n)})\,p(X \mid \theta^{(n)})} \right], then \widetilde Q(\theta^{(n)} \mid \theta^{(n)}) = \mathbb E_{p(\cdot \mid X, \theta^{(n)})} \left[ \log \frac{p(X, Z \mid \theta^{(n)})} {p(Z \mid X, \theta^{(n)})\,p(X \mid \theta^{(n)})} \right] = 0 because p(X, Z \mid \theta^{(n)}) = p(Z \mid X, \theta^{(n)})\,p(X \mid \theta^{(n)}). So at the current iterate the bound is tight, and any increase of the bound forces an increase of the observed-data log-likelihood. This is the mechanism that makes EM monotone.
Equivalently, we have the exact identity \log p(X\mid\theta)-\log p(X\mid\theta^{(n)}) =Q(\theta\mid\theta^{(n)})-Q(\theta^{(n)}\mid\theta^{(n)}) +\mathrm{KL}\left(p(\cdot\mid X,\theta^{(n)})\,\|\,p(\cdot\mid X,\theta)\right). The KL divergence is non-negative, so increasing Q increases the observed-data log-likelihood.
This gives the two EM steps:
- E-step: compute the posterior distribution p(Z \mid X, \theta^{(n)}) of the latent variables under the current parameters.
- M-step: update the parameters by maximizing \theta^{(n+1)} \in \arg\max_\theta Q(\theta \mid \theta^{(n)}).
The intuitive picture is simple: in the E-step we compute a soft guess for the missing variables, and in the M-step we pretend that these soft assignments were observed and re-fit the model. What makes EM interesting is that this heuristic picture is backed by a precise lower-bound argument, so the procedure is not just intuitive: it is guaranteed to improve the likelihood at every iteration, at least as long as each M-step really increases the surrogate objective. The EM algorithm is guaranteed to converge to a stationary point of the likelihood, but it is not guaranteed to find the global maximum if the likelihood is non-convex for instance. In practice, it is common to run EM from multiple random initializations and keep the best solution.
2. A Simple Example: Gaussian Mixture Models
Consider a Gaussian mixture model with K components. For each data point x_i, we first sample a latent component label z_i \sim \mathrm{Categorical}(\pi_1, \ldots, \pi_K), and then x_i \mid z_i = k \sim \mathcal N(\mu_k, \Sigma_k).
The parameters are \theta = \{\pi_k, \mu_k, \Sigma_k\}_{k=1}^K, and the latent variables are the component assignments Z = \{z_i\}_{i=1}^N.
For this model, the E-step computes the posterior responsibility of component k for point x_i: \gamma_{ik}^{(n)} \coloneqq p(z_i = k \mid x_i, \theta^{(n)}) = \frac{\pi_k^{(n)} \mathcal N(x_i \mid \mu_k^{(n)}, \Sigma_k^{(n)})} {\sum_{j=1}^K \pi_j^{(n)} \mathcal N(x_i \mid \mu_j^{(n)}, \Sigma_j^{(n)})}.
The corresponding EM objective is Q(\theta \mid \theta^{(n)}) = \sum_{i=1}^N \sum_{k=1}^K \gamma_{ik}^{(n)} \Big[\log \pi_k + \log \mathcal N(x_i \mid \mu_k, \Sigma_k)\Big].
Maximizing this in the M-step gives the classical updates \pi_k^{(n+1)} = \frac{1}{N}\sum_{i=1}^N \gamma_{ik}^{(n)}, \mu_k^{(n+1)} = \frac{\sum_{i=1}^N \gamma_{ik}^{(n)} x_i}{\sum_{i=1}^N \gamma_{ik}^{(n)}}, and \Sigma_k^{(n+1)} = \frac{\sum_{i=1}^N \gamma_{ik}^{(n)} (x_i - \mu_k^{(n+1)})(x_i - \mu_k^{(n+1)})^T} {\sum_{i=1}^N \gamma_{ik}^{(n)}}.
The responsibilities \gamma_{ik}^{(n)} can be read as soft cluster memberships. EM alternates between computing these soft assignments and updating each Gaussian using the weighted empirical mean and covariance of the data assigned to it.
3. A cooler example: the Baum-Welch algorithm for Speech Recognition
Another classical EM example is the hidden Markov model (HMM). Here we observe a sequence (Y_t)_{t=1}^T and introduce a latent Markov chain (X_t)_{t=1}^T. The model assumes that
- (X_t) is a Markov chain;
- conditioned on X_t, the observation Y_t is independent of the rest of the trajectory.
For simplicity, let us consider a discrete HMM. Assume X_t \in \{1, \ldots, K\}, \qquad Y_t \in \{1, \ldots, M\}. The parameters are \theta = (\pi, A, B), where \pi_i = \mathbb P(X_1 = i), \qquad A_{ij} = \mathbb P(X_{t+1} = j \mid X_t = i), \qquad B_{im} = \mathbb P(Y_t = m \mid X_t = i).
For a fixed hidden path x_{1:T} and observations y_{1:T}, the complete-data likelihood is p(y_{1:T}, x_{1:T} \mid \theta) = \pi_{x_1} \left(\prod_{t=1}^{T-1} A_{x_t x_{t+1}}\right) \left(\prod_{t=1}^{T} B_{x_t y_t}\right).
Since the joint law only involves interactions between consecutive time steps, the E-step only needs to compute the posterior marginals \gamma_t(i) \coloneqq \mathbb P(X_t = i \mid Y_{1:T} = y_{1:T}, \theta^{(n)}) and the pairwise posteriors \xi_t(i,j) \coloneqq \mathbb P(X_t = i, X_{t+1} = j \mid Y_{1:T} = y_{1:T}, \theta^{(n)}). These quantities can be obtained using the forward-backward algorithm (see Appendix).
With this notation, the EM objective becomes Q(\theta \mid \theta^{(n)}) = \sum_{i=1}^K \gamma_1(i)\log \pi_i + \sum_{t=1}^{T-1} \sum_{i=1}^K \sum_{j=1}^K \xi_t(i,j)\log A_{ij} + \sum_{t=1}^{T} \sum_{i=1}^K \gamma_t(i)\log B_{i y_t}.
The M-step is now a constrained optimization problem. We want to maximize Q(\theta \mid \theta^{(n)}) under the constraints \sum_{i=1}^K \pi_i = 1, \qquad \sum_{j=1}^K A_{ij} = 1 \text{ for each } i, \qquad \sum_{m=1}^M B_{im} = 1 \text{ for each } i. This can be done using Lagrange multipliers. We will give the derivation in the Appendix. The resulting updates are \pi_i^{(n+1)} = \gamma_1(i), A_{ij}^{(n+1)} = \frac{\sum_{t=1}^{T-1} \xi_t(i,j)} {\sum_{t=1}^{T-1} \gamma_t(i)}, and B_{im}^{(n+1)} = \frac{\sum_{t=1}^{T} \gamma_t(i)\,\mathbf 1\{y_t = m\}} {\sum_{t=1}^{T} \gamma_t(i)}.
For a corpus of independent sequences, we sum the numerators and denominators over the sequences. The update of \pi_i is then the average of \gamma_1(i) over the sequences.
In words: \gamma_t(i) is the posterior probability, or expected indicator, that state i is occupied at time t, while \xi_t(i,j) is the expected indicator of a transition from i to j between times t and t+1. Their sums over time give the expected state and transition counts. Baum-Welch simply re-estimates the initial distribution, transition matrix, and emission matrix from these soft counts.
In speech recognition, the observed sequence Y_t represents acoustic features extracted from audio, while the hidden states X_t correspond to latent speech units such as phonemes or sub-phonetic states. Real systems usually use continuous emissions rather than a small discrete matrix B, but the EM logic is exactly the same.
To make this more concrete, the figure below runs a small Baum-Welch implementation on a toy speech-recognition problem. The hidden states are three phoneme-like classes (SIL, AA, S), and the observations are coarse acoustic frames such as quiet, vowel-like, fricative-like, and burst-like slices. This is much smaller than a real acoustic model, but the mechanics are the same: the E-step computes soft alignments of hidden states to frames, and the M-step turns those soft counts into updated transition and emission probabilities.
The animation shows one focused utterance while the parameter update is computed from a small synthetic corpus. The heatmap is the posterior \gamma_t(i) = \mathbb P(X_t = i \mid Y_{1:T}, \theta^{(n)}), and the matrices on the right are the normalized soft counts that produce the next parameters \theta^{(n+1)}.
Appendix
Forward-backward derivation
We suppress the dependence on the current parameter \theta^{(n)} to keep the notation light. Define the forward messages \alpha_t(i) \coloneqq p(Y_{1:t}=y_{1:t},X_t=i\mid\theta^{(n)}). At the first time step, \alpha_1(i)=\pi_i B_{i y_1}. For t\geq 1, we condition on the previous hidden state and use the Markov property: \begin{align*} \alpha_{t+1}(j) &=p(Y_{1:t+1}=y_{1:t+1},X_{t+1}=j\mid\theta^{(n)})\\ &=B_{j y_{t+1}}\sum_{i=1}^K \alpha_t(i)A_{ij}. \end{align*}
Similarly, define the backward messages \beta_t(i)\coloneqq p(Y_{t+1:T}=y_{t+1:T}\mid X_t=i,\theta^{(n)}). Since there are no observations after time T, \beta_T(i)=1. Working backward gives \beta_t(i)=\sum_{j=1}^K A_{ij}B_{j y_{t+1}}\beta_{t+1}(j).
The likelihood of the observed sequence can be recovered at any time step: p(y_{1:T}\mid\theta^{(n)}) =\sum_{i=1}^K \alpha_t(i)\beta_t(i) =\sum_{i=1}^K \alpha_T(i). Therefore, the one-time posterior marginals are \gamma_t(i) =\frac{\alpha_t(i)\beta_t(i)} {\sum_{k=1}^K\alpha_t(k)\beta_t(k)}, and the two-time posterior marginals are \xi_t(i,j) =\frac{\alpha_t(i)A_{ij}B_{j y_{t+1}}\beta_{t+1}(j)} {\sum_{k=1}^K\sum_{\ell=1}^K \alpha_t(k)A_{k\ell}B_{\ell y_{t+1}}\beta_{t+1}(\ell)}. The algorithm costs O(TK^2) operations, instead of summing over all K^T hidden paths.
In practice, the forward messages quickly become numerically very small. We therefore normalize them at every time step and divide the backward messages by the same normalization constants. This does not change \gamma_t or \xi_t, and the log-likelihood is the sum of the logarithms of the forward normalization constants.
Lagrange multiplier derivation of the HMM M-step
The terms of Q involving \pi, A, and B can be optimized separately. Let C_{ij}\coloneqq\sum_{t=1}^{T-1}\xi_t(i,j), \qquad D_{im}\coloneqq\sum_{t=1}^T\gamma_t(i)\mathbf 1\{y_t=m\}.
For the initial distribution, introduce a Lagrange multiplier for the constraint \sum_i\pi_i=1: \mathcal J_\pi =\sum_{i=1}^K\gamma_1(i)\log\pi_i +\lambda\left(\sum_{i=1}^K\pi_i-1\right). The stationarity condition is \frac{\partial\mathcal J_\pi}{\partial\pi_i} =\frac{\gamma_1(i)}{\pi_i}+\lambda=0. Thus \pi_i is proportional to \gamma_1(i). Since both distributions sum to one, \pi_i^{(n+1)}=\gamma_1(i).
For a fixed state i, the part of Q involving the i-th row of A is \sum_j C_{ij}\log A_{ij}. Introducing one multiplier \lambda_i for this row gives \mathcal J_{A_i} =\sum_{j=1}^K C_{ij}\log A_{ij} +\lambda_i\left(\sum_{j=1}^K A_{ij}-1\right). The stationarity condition shows that A_{ij} is proportional to C_{ij}, hence A_{ij}^{(n+1)} =\frac{C_{ij}}{\sum_{\ell=1}^K C_{i\ell}} =\frac{\sum_{t=1}^{T-1}\xi_t(i,j)} {\sum_{t=1}^{T-1}\gamma_t(i)}, where we used \sum_j\xi_t(i,j)=\gamma_t(i).
The emission matrix is identical. For a fixed state i, maximize \sum_m D_{im}\log B_{im} under \sum_m B_{im}=1. This gives B_{im}^{(n+1)} =\frac{D_{im}}{\sum_{\ell=1}^M D_{i\ell}} =\frac{\sum_{t=1}^T\gamma_t(i)\mathbf 1\{y_t=m\}} {\sum_{t=1}^T\gamma_t(i)}.
These formulas assume that the denominators are nonzero. If a state receives no posterior mass, its transition or emission row is not identified by the current E-step; in practice one can leave that row unchanged or add a small pseudocount.