EM algorithm

Topic of Mixtures and the EM algorithm, part of Statistical learning · Also in Mixtures and the EM algorithm

EM algorithm

Type Topic
Taught in Mixtures and the EM algorithm · Mixtures and the EM algorithm
Part Statistical learning

An iterative method for maximum likelihood with missing or latent data. The E step computes the expected complete-data log-likelihood given the current parameters; the M step maximises it. No iteration decreases the observed-data likelihood, but the limit can be a local maximum, so several starts are used.

In Mixtures and the EM algorithm

Let \Psi^{(k)} be the parameter at iteration k. The update maximises the conditional expected complete log-likelihood. The expectation is over the latent indicators, drawn from their conditional distribution at \Psi^{(k)}, with the observations held fixed. That expectation is the Q function,

Q(\Psi;\Psi^{(k)}) =\mathbb{E}_{Z;\Psi^{(k)}}\bigl[\ell_C(\Psi)\bigm|X_{\mathrm{all}}\bigr].

The two arguments of Q have different jobs. The current value \Psi^{(k)} determines the distribution of the missing labels. The candidate \Psi is what varies during maximisation. Compute the memberships once in the E step, hold them fixed in the M step, then recompute them at the new parameter. Updating the memberships inside the M-step derivatives would be a different calculation.

The E step

Main article: E step

The complete log-likelihood is linear in the indicators, so the conditional expectation passes inside the sum. The expectation of a binary indicator is the probability that it equals one. Bayes’ rule for a discrete label and a continuous measurement identifies that probability with the posterior membership, evaluated at the current parameter:

\begin{align*} Q(\Psi;\Psi^{(k)}) &=\mathbb{E}\Biggl[ \sum_{i=1}^{n}\sum_{h=1}^{H}Z_{ih}\bigl[\ln\pi_h+\ln f_h(\mathbf{x}_i;\theta_h)\bigr] \Bigm|X_{\mathrm{all}} \Biggr]\\ &=\sum_{i=1}^{n}\sum_{h=1}^{H} \mathbb{E}(Z_{ih}\mid\mathbf{x}_i)\, \bigl[\ln\pi_h+\ln f_h(\mathbf{x}_i;\theta_h)\bigr]\\ &=\sum_{i=1}^{n}\sum_{h=1}^{H} \tau_h(\mathbf{x}_i;\Psi^{(k)})\, \bigl[\ln\pi_h+\ln f_h(\mathbf{x}_i;\theta_h)\bigr]. \end{align*}

The current parameter \Psi^{(k)} is held fixed inside \tau_h. The free parameter \Psi remains in \ln\pi_h+\ln f_h, which is what the M step maximises. Explicitly,

\tau_h\bigl(\mathbf{x}_i;\Psi^{(k)}\bigr) =\frac{\pi_h^{(k)}f_h\bigl(\mathbf{x}_i;\theta_h^{(k)}\bigr)} {\sum_{m=1}^{H}\pi_m^{(k)}f_m\bigl(\mathbf{x}_i;\theta_m^{(k)}\bigr)}.

The M step

Main article: M step

Maximise Q in \Psi, with the weights \tau_{ih}:=\tau_h(\mathbf{x}_i;\Psi^{(k)}) held at the values just computed. The objective separates into a term in the weights \pi_h and a weighted normal log-likelihood for each component. Write N_h=\sum_i\tau_{ih} for the effective sample size of component h.

Weights. Maximise \sum_h N_h\ln\pi_h subject to \sum_h\pi_h=1. The Lagrangian \sum_h N_h\ln\pi_h-\kappa\bigl(\sum_h\pi_h-1\bigr) has derivative N_h/\pi_h-\kappa. Set the derivative to zero, so \pi_h=N_h/\kappa. Sum over h. The identity \sum_h N_h=\sum_i\sum_h\tau_{ih}=n forces \kappa=n, and therefore

\pi_h^{(k+1)} =\frac1n\sum_{i=1}^{n}\tau_h\bigl(\mathbf{x}_i;\Psi^{(k)}\bigr) =\frac{N_h}{n}.

Means. The terms of Q that involve \boldsymbol\mu_h are -\tfrac12\sum_i\tau_{ih}(\mathbf{x}_i-\boldsymbol\mu_h)^\top\Sigma_h^{-1}(\mathbf{x}_i-\boldsymbol\mu_h). The gradient in \boldsymbol\mu_h is \Sigma_h^{-1}\sum_i\tau_{ih}(\mathbf{x}_i-\boldsymbol\mu_h). Set it to zero and multiply by \Sigma_h, which is positive definite at an interior point. The solution is the weighted mean,

\boldsymbol\mu_h^{(k+1)} =\frac{\sum_{i=1}^{n}\tau_{ih}\mathbf{x}_i}{\sum_{i=1}^{n}\tau_{ih}}.

Covariance matrices. Pass to the precision \Theta_h=\Sigma_h^{-1}. The terms that involve \Theta_h are \tfrac12 N_h\ln|\Theta_h|-\tfrac12\sum_i\tau_{ih}(\mathbf{x}_i-\boldsymbol\mu_h)^\top\Theta_h(\mathbf{x}_i-\boldsymbol\mu_h). The matrix derivatives \partial\ln|\Theta|/\partial\Theta=\Theta^{-1} and \partial(\mathbf{v}^\top\Theta\mathbf{v})/\partial\Theta=\mathbf{v}\mathbf{v}^\top produce

\frac12 N_h\Theta_h^{-1} -\frac12\sum_{i=1}^{n}\tau_{ih}(\mathbf{x}_i-\boldsymbol\mu_h)(\mathbf{x}_i-\boldsymbol\mu_h)^\top.

Set the derivative to zero and insert the updated mean. Then \Theta_h^{-1} equals the weighted matrix of squares and products, so

\Sigma_h^{(k+1)} =\frac{\sum_{i=1}^{n}\tau_{ih} \bigl(\mathbf{x}_i-\boldsymbol\mu_h^{(k+1)}\bigr) \bigl(\mathbf{x}_i-\boldsymbol\mu_h^{(k+1)}\bigr)^\top} {\sum_{i=1}^{n}\tau_{ih}}.

The formulas now have a familiar shape: proportions, means, and covariance matrices, with fractional membership replacing a hard group label. The divisor is N_h. Likelihood maximisation uses N_h in the denominator; an unbiased sample covariance would subtract one from a hard group size, and that correction is a different estimator. These updates require N_h>0 and positive definite updated covariance matrices. A collapsed component is a failure of those conditions.

Definition 1 (EM for a normal mixture) Choose initial values \Psi^{(0)}. For k=0,1,2,\ldots:

  1. E step: compute \tau_{ih}=\tau_h(\mathbf{x}_i;\Psi^{(k)}) for all i,h.
  2. M step: compute \pi_h^{(k+1)}, \boldsymbol\mu_h^{(k+1)}, and \Sigma_h^{(k+1)} from the three updates above.
  3. Stop when the change in \ell(\Psi^{(k)};X_{\mathrm{all}}) is smaller than a threshold.

Article: EM for a normal mixture →

Q(\Psi;\Psi^{(k)}), plus a constant independent of \Psi, is a lower bound on the observed log-likelihood, touching it at the current parameter. The constant matters for the touching statement, and it does not change the maximiser. EM is the minorisation–maximisation step for that bound (Lange, 2013).

In Mixtures and the EM algorithm, continued

Theorem 1 (Monotonicity of EM) An exact EM step cannot decrease the observed log-likelihood (Dempster, Laird, and Rubin, 1977). Convergence of the likelihood values requires an upper bound. Convergence of the parameters requires further conditions. Under the regularity conditions studied by Wu (1983), limit points are stationary. A stationary point of a mixture likelihood can fail to be a global maximum.

Article: Monotonicity of EM →

Write the observed log-density by splitting off the conditional probability of the latent labels,

\ln f(X_{\mathrm{all}};\Psi) =\ell_C(\Psi;X_{\mathrm{all}},Z)-\ln\mathbb{P}(Z\mid X_{\mathrm{all}};\Psi).

Take the expectation over Z drawn from the conditional distribution at the current parameter \Psi^{(k)}. The first term becomes Q, and the second becomes a function G:

\ell(\Psi)=Q(\Psi;\Psi^{(k)})-G(\Psi;\Psi^{(k)}), \qquad G(\Psi;\Psi^{(k)}) =\mathbb{E}_{\Psi^{(k)}}\bigl[\ln\mathbb{P}(Z\mid X_{\mathrm{all}};\Psi)\bigm|X_{\mathrm{all}}\bigr].

Gibbs’ inequality is nonnegativity of Kullback–Leibler divergence, so G(\Psi;\Psi^{(k)})\le G(\Psi^{(k)};\Psi^{(k)}) for every \Psi. Split the change in the observed log-likelihood:

\begin{align*} \ell(\Psi^{(k+1)})-\ell(\Psi^{(k)}) &= \underbrace{Q(\Psi^{(k+1)};\Psi^{(k)})-Q(\Psi^{(k)};\Psi^{(k)})}_{\ge 0\text{ (M step)}}\\ &\qquad+ \underbrace{G(\Psi^{(k)};\Psi^{(k)})-G(\Psi^{(k+1)};\Psi^{(k)})}_{\ge 0\text{ (Gibbs)}}. \end{align*}

Both terms are nonnegative. Increasing Q is enough; finding its exact maximum is sufficient and stronger than necessary. An algorithm that only increases Q is generalised EM. At a differentiable interior fixed point, with differentiation under the expectation justified, the score of G vanishes and the score of \ell equals the score of Q. Such a fixed point is stationary. The interior and smoothness qualifications matter near zero weights or singular covariances.

The lower-bound form is the same identity rearranged. Because G(\Psi;\Psi^{(k)})\le G(\Psi^{(k)};\Psi^{(k)}),

\ell(\Psi)\ge Q(\Psi;\Psi^{(k)})-G(\Psi^{(k)};\Psi^{(k)}),

with equality at \Psi=\Psi^{(k)}. The right-hand side differs from Q by a constant, so the M step maximises a surface that touches \ell at the current parameter.

Example 1 (Two iterations on six points) Question. Take the one-dimensional data x=(-1.2,-0.4,0.3,2.1,2.8,3.5) and a two-component normal mixture with starting values \pi=(0.5,0.5), \mu=(0,3), \sigma^2=(1,1). Perform the E and M steps.

Solution. E step. The component densities and responsibilities, with \varphi(\cdot;\mu,\sigma^2) the \operatorname{N}(\mu,\sigma^2) density, are

x_i -1.2 -0.4 0.3 2.1 2.8 3.5
f_1(x_i)=\varphi(x_i;0,1) 0.1942 0.3683 0.3814 0.0440 0.0079 0.0009
f_2(x_i)=\varphi(x_i;3,1) 0.0001 0.0012 0.0104 0.2661 0.3910 0.3521
\tau_{i1} 0.9997 0.9967 0.9734 0.1419 0.0198 0.0025
\tau_{i2}=1-\tau_{i1} 0.0003 0.0033 0.0266 0.8581 0.9802 0.9975

For the fourth point,

\tau_{41} =\frac{0.5\times 0.0440}{0.5\times 0.0440+0.5\times 0.2661} =\frac{0.0220}{0.15505} =0.1419.

The mixture densities 0.5f_1+0.5f_2 are 0.09715, 0.18475, 0.19590, 0.15505, 0.19945, 0.17650. Their logarithms are -2.3315, -1.6888, -1.6302, -1.8640, -1.6122, -1.7344, and

\ell=\sum_i\ln\bigl(0.5f_1(x_i)+0.5f_2(x_i)\bigr)=-10.861.

M step. The four-decimal memberships in the table sum to 3.1340 and 2.8660. The effective sizes used for the update, before that rounding, are N_1=\sum_i\tau_{i1}=3.1339 and N_2=2.8661, so

\pi^{(1)}=\Bigl(\frac{3.1339}{6},\frac{2.8661}{6}\Bigr)=(0.5223,0.4777).

The weighted sum for the first mean is

\begin{align*} \sum_i\tau_{i1}x_i &=0.9997(-1.2)+0.9967(-0.4)+0.9734(0.3)\\ &\qquad+0.1419(2.1)+0.0198(2.8)+0.0025(3.5)\\ &=-0.94412, \end{align*}

so \mu_1^{(1)}=-0.94412/3.1339=-0.3013. The same calculation for component 2 gives \mu_2^{(1)}=2.8067. The weighted variances about these new means are \sigma_1^{2(1)}=0.7063 and \sigma_2^{2(1)}=0.3888.

The log-likelihood at the next iteration is -9.808, and the one after that is -9.539. After three iterations the estimates are \pi\to(0.5002,0.4998), \mu\to(-0.4323,2.8002), \sigma^2\to(0.378,0.327). The observed log-likelihood is non-decreasing, as the argument above guarantees.

The fragment implements the normal-mixture iteration just derived. It normalises log densities after subtracting their row maximum, which avoids many underflow problems. The returned responsibilities and final likelihood are recomputed at the returned parameters. This is an educational unconstrained fit: it stops if a covariance loses positive definiteness or a component has zero numerical weight. It does not silently replace the model with a regularised one.

em_mixture <- function(x, mu, Sigma, pi, iters = 100L) {
  x <- as.matrix(x); n <- nrow(x); H <- length(pi)
  stopifnot(iters >= 0, iters == as.integer(iters),
            H == length(mu), H == length(Sigma),
            all(pi > 0), abs(sum(pi) - 1) < 1e-8)
  expectation <- function() {
    # Cholesky fails explicitly if a covariance is not positive definite.
    invisible(lapply(Sigma, chol))
    ld <- vapply(seq_len(H), function(h)
      log(pi[h]) + mvtnorm::dmvnorm(x, mu[[h]], Sigma[[h]], log = TRUE),
      numeric(n))
    a <- apply(ld, 1, max)
    stopifnot(all(is.finite(a)))
    e <- exp(ld - a) # subtract before exponentiating
    z <- rowSums(e)
    list(tau = e / z, ll = sum(a + log(z)))
  }
  state <- expectation()
  loglik <- c(state$ll, numeric(iters))
  for (it in seq_len(iters)) {
    tau <- state$tau
    Nh <- colSums(tau)
    stopifnot(all(Nh > 0))
    pi <- Nh / n
    for (h in seq_len(H)) {
      mu[[h]] <- colSums(tau[, h] * x) / Nh[h]
      xc <- sweep(x, 2, mu[[h]])
      Sigma[[h]] <- crossprod(xc * tau[, h], xc) / Nh[h]
    }
    state <- expectation() # evaluate the updated parameters
    loglik[it + 1L] <- state$ll
  }
  list(pi = pi, mu = mu, Sigma = Sigma,
       tau = state$tau, loglik = loglik)
}

Used in

See also

EM for a normal mixture · Monotonicity of EM · Complete-data likelihood · Maximum likelihood · E step · M step · Likelihood function · Gibbs’ inequality

Glossary and index