Lab 8: The EM Algorithm

STAT 536 — Statistical Computing

Author

Soumojit Das

Published

Fall 2026

Learning Objectives

By the end of today’s lab, you will:

Understand the E-step (computing expectations) and the M-step (maximization)
Implement the E-step and the M-step of EM for three problems: genetic linkage, censored survival data, and a Gaussian mixture
Compare EM with direct optimization methods: BFGS through optim(), and a closed-form MLE
Check that the log-likelihood never decreases along the EM path, from the starting value on
Start a mixture from k-means, and see what EM does when every component starts at the same values

Setup

Figure labels are typeset by LaTeX. tikz_png() draws a ggplot through tikzDevice, compiles it with pdflatex, and returns a PNG. It needs a LaTeX installation (TinyTeX is enough) and installs tikzDevice, pdftools, and png on first use.

library(ggplot2)
library(patchwork)
set.seed(536)
theme_set(theme_bw(base_size = 14))

pal <- c(red = "#D62828", blue = "#2E86AB", green = "#06A77D", orange = "#F77F00", dark = "#023047", grey = "#6a6a6a")

find_pdflatex <- function() {
  cand <- c(Sys.which("pdflatex"),
            path.expand("~/Library/TinyTeX/bin/universal-darwin/pdflatex"),
            file.path(Sys.getenv("APPDATA"), "TinyTeX/bin/windows/pdflatex.exe"),
            path.expand("~/.TinyTeX/bin/x86_64-linux/pdflatex"))
  cand <- cand[nzchar(cand) & file.exists(cand)]
  if (!length(cand)) stop("pdflatex not found; install TinyTeX with tinytex::install_tinytex()")
  cand[[1]]
}

tikz_png <- function(p, png_file, width = 9, height = 5, dpi = 150) {
  for (pkg in c("tikzDevice", "pdftools", "png"))
    if (!requireNamespace(pkg, quietly = TRUE)) install.packages(pkg)
  options(tikzLatex = find_pdflatex(),
          tikzDocumentDeclaration = r"(\documentclass[10pt]{article})",
          tikzLatexPackages = c(r"(\usepackage{tikz})",
                                r"(\usepackage[active,tightpage,psfixbb]{preview})",
                                r"(\PreviewEnvironment{pgfpicture})",
                                r"(\setlength\PreviewBorder{0pt})",
                                r"(\usepackage[T1]{fontenc})", r"(\usepackage{lmodern})",
                                r"(\usepackage{amsmath})"))
  work <- file.path(tempdir(), "tikz_png"); dir.create(work, showWarnings = FALSE)
  tex <- file.path(work, "fig.tex"); pdf <- file.path(work, "fig.pdf")
  tikzDevice::tikz(tex, width = width, height = height, standAlone = TRUE)
  print(p); dev.off()
  ok <- system2(getOption("tikzLatex"),
                c("-interaction=batchmode", "-halt-on-error",
                  paste0("-output-directory=", shQuote(work)), shQuote(tex)),
                stdout = FALSE, stderr = FALSE)
  if (ok != 0 || !file.exists(pdf)) stop("pdflatex failed; see ", file.path(work, "fig.log"))
  dir.create(dirname(png_file), recursive = TRUE, showWarnings = FALSE)
  png::writePNG(pdftools::pdf_render_page(pdf, page = 1, dpi = dpi), png_file)
  png_file
}

# tick labels in math mode, so minus signs and decimals are set by LaTeX
tex_num <- function(v) ifelse(is.na(v), NA,
  paste0("$", format(v, scientific = FALSE, drop0trailing = TRUE, trim = TRUE), "$"))
tex_theme <- theme(text = element_text(size = 17), plot.title = element_text(hjust = 0.5),
                   legend.position = "top", legend.key.width = unit(1.2, "cm"))
# whole-number ticks for an iteration axis
int_breaks <- function(lim) unique(round(pretty(lim)))

# Utility function for convergence plots
# (entry 1 of log_lik is iteration 0, the starting value)
plot_convergence <- function(log_lik, title = "EM algorithm convergence") {
  df <- data.frame(
    iteration = seq_along(log_lik) - 1,
    log_likelihood = log_lik
  )

  ggplot(df, aes(x = iteration, y = log_likelihood)) +
    geom_line(color = pal[["blue"]], linewidth = 1) +
    geom_point(color = pal[["red"]], size = 2) +
    scale_x_continuous(breaks = int_breaks, labels = tex_num) +
    scale_y_continuous(labels = tex_num) +
    labs(title = title,
         x = "Iteration",
         y = "Log-likelihood") +
    tex_theme
}

The EM Algorithm

Recall the EM algorithm from lecture? The notation is the lecture’s.

Symbol	Meaning
$X$; $Z$	the observed data; the latent (missing) data
$\theta$	the parameter ($\psi$ in Part 1)
$\theta^{(t)}$	the parameter after $t$ iterations; $\theta^{(0)}$ is the starting value
$\ell(\theta; X)$	the observed-data log-likelihood
$\ell_c(\theta; X, Z) = \log p(X, Z \mid \theta)$	the complete-data log-likelihood
$p(Z \mid X, \theta^{(t)})$	the posterior of $Z$ at $\theta^{(t)}$
$Q(\theta \mid \theta^{(t)})$	the expected complete-data log-likelihood

EM Algorithm

Input: Observed data $X$, initial parameter $\theta^{(0)}$

Output: The converged value $\hat\theta$, a stationary point of $\ell$ (keep the best of several starts)

Repeat for $t = 0, 1, 2, \ldots$ until convergence:

E-step: Compute the expected complete-data log-likelihood: \[ \begin{aligned} Q(\theta \mid \theta^{(t)}) &= \mathbb{E}_{Z \mid X, \theta^{(t)}}[\ell_c(\theta; X, Z)] \\ &= \mathbb{E}_{Z \mid X, \theta^{(t)}}[\log p(X, Z \mid \theta)] = \int p(Z \mid X, \theta^{(t)}) \log p(X, Z \mid \theta)\,dZ \end{aligned} \]

M-step: Update parameter by maximizing $Q$: $\theta^{(t+1)} = \argmax_\theta Q(\theta \mid \theta^{(t)})$

Return $\hat\theta = \theta^{(t+1)}$

Essentially, this says that EM averages the complete-data log-likelihood over the missing $Z$, using the posterior of $Z$ at the current parameter, and then maximizes the average. By the monotonicity theorem from lecture, no EM iteration can lower the likelihood: $\ell(\theta^{(t+1)}; X) \ge \ell(\theta^{(t)}; X)$.

In all three parts of this lab the terms of $\ell_c$ that involve the parameter are linear in the missing data. The E-step then reduces to the conditional expectation of the missing data.

Part 1: Genetic Linkage

This is the classic example from Dempster, Laird, and Rubin (1977). We observe counts of genetic phenotypes but don’t observe the underlying genotypes.

Background: Consider a genetic model with four observable phenotypes. The counts follow a multinomial distribution, but one of the cell counts is a mixture of two unobservable genotype categories.

Observed Data: $X = (X_1, X_2, X_3, X_4) = (125, 18, 20, 34)$

Model: The cell probabilities are: \[p_1 = \frac{1}{2} + \frac{\psi}{4}, \quad p_2 = \frac{1 - \psi}{4}, \quad p_3 = \frac{1 - \psi}{4}, \quad p_4 = \frac{\psi}{4}\]

where $p_j$ is the probability that an observation falls in cell $j$, and $\psi \in [0, 1]$ is the parameter we want to estimate.

The Challenge: Cell 1 is a mixture. If we knew how many came from each component, the MLE would be trivial.

Complete Data: Imagine we could split $X_1$ into:

$Z_1$: count from the $1/2$ component
$Z_2$: count from the $\psi/4$ component
So $X_1 = Z_1 + Z_2$

With the split counts, the complete-data log-likelihood is \[\ell_c(\psi; X, Z) = (Z_2 + X_4)\log\psi + (X_2 + X_3)\log(1 - \psi) + \text{terms without } \psi,\] which is linear in the missing count $Z_2$.

set.seed(536)
# Observed data
x <- c(125, 18, 20, 34)
n <- sum(x)

print("Observed counts:")
print(sprintf("X1 = %d (mixture cell)", x[1]))
print(sprintf("X2 = %d", x[2]))
print(sprintf("X3 = %d", x[3]))
print(sprintf("X4 = %d", x[4]))
print(sprintf("Total n = %d", n))

E-Step: Given current estimate $\psi^{(t)}$, compute the expected value of $Z_2$ (the “missing” part of $X_1$):

\[\E_{Z_2 \mid X_1, \psi^{(t)}}[Z_2] = X_1 \cdot \frac{\psi^{(t)}/4}{1/2 + \psi^{(t)}/4}\]

This is just the posterior probability that an observation in cell 1 came from the $\psi/4$ component, multiplied by $X_1$.

M-Step: If we knew $Z_2$, the MLE would be: \[\hat{\psi} = \frac{Z_2 + X_4}{Z_2 + X_2 + X_3 + X_4}\]

Since we don’t know $Z_2$, we use its expected value from the E-step.

Log-Likelihood: The observed log-likelihood is: \[ \begin{aligned} \ell(\psi) = {}& X_1 \log\left(\frac{1}{2} + \frac{\psi}{4}\right) + X_2 \log\left(\frac{1-\psi}{4}\right) \\ &+ X_3 \log\left(\frac{1-\psi}{4}\right) + X_4 \log\left(\frac{\psi}{4}\right) \end{aligned} \]

This leaves out the logarithm of the multinomial coefficient, a constant that does not involve $\psi$.

Exercise 1.1: Implementation

Complete the EM algorithm function below: the E-step and the M-step.
Run it from $\psi^{(0)} = 0.5$.

set.seed(536)
# EM algorithm for genetic linkage
genetic_em <- function(x, psi_init = 0.5, max_iter = 100, tol = 1e-8) {
  # Observed-data log-likelihood
  obs_log_lik <- function(psi) {
    p1 <- 1/2 + psi/4
    p2 <- (1 - psi)/4
    p3 <- (1 - psi)/4
    p4 <- psi/4
    x[1] * log(p1) + x[2] * log(p2) + x[3] * log(p3) + x[4] * log(p4)
  }
  psi <- psi_init
  # entry 1 holds iteration 0, the starting value
  log_lik <- c(obs_log_lik(psi), numeric(max_iter))
  psi_history <- c(psi, numeric(max_iter))

  for (iter in 1:max_iter) {
    psi_old <- psi

    # TODO: E-step. Compute the expected value of Z2 given current psi.
    # Hint: prob_z2 = (psi / 4) / (1/2 + psi / 4)
    #       E_z2 = x[1] * prob_z2
    prob_z2 <- ___________
    E_z2 <- ___________

    # TODO: M-step. Update psi using the expected Z2.
    # Hint: psi = (E_z2 + x[4]) / (E_z2 + x[2] + x[3] + x[4])
    psi <- ___________

    # Compute log-likelihood
    log_lik[iter + 1] <- obs_log_lik(psi)
    psi_history[iter + 1] <- psi

    # Check convergence
    if (abs(psi - psi_old) < tol) {
      log_lik <- log_lik[1:(iter + 1)]
      psi_history <- psi_history[1:(iter + 1)]
      break
    }
  }

  list(psi = psi,
       iterations = length(log_lik) - 1,
       log_lik = log_lik,
       psi_history = psi_history)
}

# Run EM algorithm
result_em <- genetic_em(x, psi_init = 0.5)

print("EM Algorithm Results:")
print(sprintf("Estimated psi: %.6f", result_em$psi))
print(sprintf("Converged in %d iterations", result_em$iterations))
print(sprintf("Final log-likelihood: %.4f", tail(result_em$log_lik, 1)))

Report.

The estimate of $\psi$.
The number of iterations.
The final log-likelihood.

Exercise 1.2: Comparison with Direct Optimization

Let’s compare EM with a quasi-Newton method (BFGS, through optim).

The code chunk below is complete. Run it as is. It maximizes $\ell(\psi)$ directly. optim() minimizes, so it is given $-\ell(\psi)$.

set.seed(536)
# Direct optimization of observed likelihood
log_likelihood <- function(psi, x) {
  if (psi < 0 || psi > 1) return(-Inf)
  p1 <- 1/2 + psi/4
  p2 <- (1 - psi)/4
  p3 <- (1 - psi)/4
  p4 <- psi/4
  x[1] * log(p1) + x[2] * log(p2) + x[3] * log(p3) + x[4] * log(p4)
}

# Use optim with BFGS (quasi-Newton)
result_optim <- optim(
  par = 0.5,
  fn = function(psi) -log_likelihood(psi, x),  # negative for minimization
  method = "BFGS",
  control = list(maxit = 100)
)

print("Quasi-Newton (BFGS) Results:")
print(sprintf("Estimated psi: %.6f", result_optim$par))
print(sprintf("Function evaluations: %d", result_optim$counts[1]))
print(sprintf("Gradient evaluations (finite differences): %d",
              result_optim$counts[2]))
print(sprintf("Final log-likelihood: %.4f", -result_optim$value))

print("--- COMPARISON ---")
print(sprintf("EM estimate:     %.6f (%d iterations)",
              result_em$psi, result_em$iterations))
print(sprintf("BFGS estimate:   %.6f (%d gradient evaluations)",
              result_optim$par, result_optim$counts[2]))
print(sprintf("Difference:      %.8f",
              abs(result_em$psi - result_optim$par)))

Report.

The BFGS estimate of $\psi$ and its final log-likelihood.
The numbers of function and gradient evaluations.
The difference between the EM and BFGS estimates.

Exercise 1.3: Visualizing Convergence

Run the figure chunk. The left panel is the observed log-likelihood $\ell(\psi^{(t)})$ at each iteration $t$, from the starting value at $t = 0$. The right panel is $\psi^{(t)}$. The dashed line is the final estimate.

set.seed(536)
# Plot convergence
p1 <- plot_convergence(result_em$log_lik, "Log-likelihood convergence")

# Plot parameter trajectory
df_trajectory <- data.frame(
  iteration = 0:result_em$iterations,
  psi = result_em$psi_history
)

p2 <- ggplot(df_trajectory, aes(x = iteration, y = psi)) +
  geom_line(color = pal[["blue"]], linewidth = 1) +
  geom_point(color = pal[["red"]], size = 2) +
  geom_hline(yintercept = result_em$psi,
             linetype = "dashed", color = pal[["green"]]) +
  annotate("text", x = result_em$iterations * 0.7,
           y = result_em$psi + 0.01,
           label = sprintf("Final: %.4f", result_em$psi),
           color = pal[["green"]], size = 5.5) +
  scale_x_continuous(breaks = int_breaks, labels = tex_num) +
  scale_y_continuous(labels = tex_num) +
  labs(title = "Parameter trajectory",
       x = "Iteration",
       y = r"($\psi$)") +
  tex_theme

# Combine plots
p <- p1 + p2
knitr::include_graphics(tikz_png(p, "figures/lab08/linkage-convergence.png",
                                 width = 10, height = 4.2))

Report.

The figure.
The increase in the log-likelihood at each of the first three iterations, diff(result_em$log_lik)[1:3], and the smallest increase over all iterations.

Part 2: Right-Censored Exponential Data

What is censoring?

In survival analysis and reliability studies, we often cannot observe exact event times for all subjects. This is called censoring.

Real-World Example: Clinical Trial

Imagine a cancer study measuring time until disease progression:

Patient A: Progresses at 8 months $\to$ We observe the exact time (8 months)
Patient B: Still healthy when study ends at 12 months $\to$ Right-censored at 12 months
Patient C: Moves away at 6 months, still healthy $\to$ Right-censored at 6 months

For censored patients, we know: “The event time is at least X months” but not the exact time.

Why this matters.

Ignoring censored observations wastes data and can severely bias estimates. EM provides a solution.

Types of Censoring:

Right censoring (today’s focus): We know $T > C$ (event after censoring time)
Left censoring: We know $T < C$ (event before observation began)
Interval censoring: We know $L < T < R$ (event in an interval)

True Data (unobserved for censored cases): \[T_i \stackrel{iid}{\sim} \text{Exp}(\theta) \quad \text{for } i = 1, \ldots, n\]

The exponential distribution has density: \[f(t | \theta) = \theta e^{-\theta t}, \quad t \geq 0\]

with mean $E[T] = 1/\theta$ and survival function $S(t) = P(T > t) = e^{-\theta t}$.

Observed Data:

$Y_i = \min(T_i, C_i)$ (observed time: either event or censoring time)
$\Delta_i = I(T_i \leq C_i)$ (event indicator: 1 if event observed, 0 if censored)

Here $C_i$ is the censoring time of subject $i$, independent of $T_i$.

Why this is hard.

For censored observations ($\Delta_i = 0$), we don’t know the true event time $T_i$, only that $T_i > Y_i$.

An observed event ($\Delta_i = 1$) contributes the density at $Y_i$, $\theta e^{-\theta Y_i}$. A censored time ($\Delta_i = 0$) tells us only that $T_i > Y_i$, so it contributes the survival probability $S(Y_i) = e^{-\theta Y_i}$.

Observed Likelihood: \[L(\theta; Y, \Delta) = \prod_{i=1}^n [\theta e^{-\theta Y_i}]^{\Delta_i} [e^{-\theta Y_i}]^{1 - \Delta_i} = \theta^{\sum \Delta_i} e^{-\theta \sum Y_i}\]

This is actually easy to maximize directly. But let’s see how EM handles it.

Complete Data: If we knew all $T_i$, the MLE would be trivial: \[\hat{\theta} = \frac{n}{\sum_{i=1}^n T_i}\]

It maximizes the complete-data log-likelihood $\ell_c(\theta; T) = n\log\theta - \theta\sum_{i=1}^n T_i$, which is linear in the $T_i$.

Idea: Use EM to “fill in” the missing $T_i$ for censored observations.

E-Step: For censored observations, compute: \[\E_{T_i \mid T_i > Y_i,\, \theta^{(t)}}[T_i] = Y_i + \frac{1}{\theta^{(t)}}\]

This uses the memoryless property of the exponential distribution.

Intuition: If someone has survived to time $Y_i$, the expected additional survival time is still $1/\theta$ (memoryless), so total expected time is $Y_i + 1/\theta$.

M-Step: Update parameter using expected complete data: \[\theta^{(t+1)} = \frac{n}{\sum_{i=1}^n \E_{T_i \mid Y_i, \Delta_i, \theta^{(t)}}[T_i]}\]

Exercise 2.1: Generate Censored Data

The code chunk below is complete. Run it as is. It draws $n = 100$ event times from $\text{Exp}(0.5)$ and, independently, 100 censoring times from $\text{Exp}(0.3)$.

set.seed(536)
# Generate complete survival data
n <- 100
true_theta <- 0.5  # True rate parameter (mean survival = 1/0.5 = 2)
true_times <- rexp(n, rate = true_theta)

# Generate censoring times (independent exponential with different rate)
censor_rate <- 0.3
censor_times <- rexp(n, rate = censor_rate)

# Create observed data
Y <- pmin(true_times, censor_times)
Delta <- as.numeric(true_times <= censor_times)

censoring_pct <- 100 * (1 - mean(Delta))

print("Censored Data Summary:")
print(sprintf("Sample size: %d", n))
print(sprintf("True theta: %.3f (true mean survival: %.2f)",
              true_theta, 1/true_theta))
print(sprintf("Events observed: %d (%.1f%%)",
              sum(Delta), 100 * mean(Delta)))
print(sprintf("Censored: %d (%.1f%%)", sum(1 - Delta), censoring_pct))
print(sprintf("Mean observed time: %.3f", mean(Y)))

Report.

The numbers of events and of censored observations.
The mean observed time.

Exercise 2.2: Visualize the Censoring

Run the figure chunk. The left panel shows the 50 shortest observed times, one line per subject from time 0 to $Y_i$. The right panel is a histogram of all 100 observed times. The colors mark events and censored observations.

set.seed(536)
# Create data frame for visualization
df_censor <- data.frame(
  id = 1:n,
  time = Y,
  status = factor(Delta, levels = c(0, 1), labels = c("Censored", "Event"))
)

# Sort by time for better visualization
df_censor <- df_censor[order(df_censor$time), ]
df_censor$id <- 1:n

p1 <- ggplot(df_censor[1:50, ], aes(x = time, y = id, color = status)) +
  geom_point(size = 2) +
  geom_segment(aes(x = 0, xend = time, y = id, yend = id), alpha = 0.3) +
  scale_color_manual(values = c("Censored" = pal[["red"]],
                                "Event" = pal[["blue"]])) +
  scale_x_continuous(labels = tex_num) +
  scale_y_continuous(labels = tex_num) +
  labs(title = "The 50 shortest observed times", color = NULL,
       x = "Observed time",
       y = "Subject, sorted by time") +
  tex_theme

# Histogram of observed times by status
p2 <- ggplot(df_censor, aes(x = time, fill = status)) +
  geom_histogram(bins = 30, alpha = 0.7, position = "stack") +
  scale_fill_manual(values = c("Censored" = pal[["red"]],
                               "Event" = pal[["blue"]])) +
  scale_x_continuous(labels = tex_num) +
  scale_y_continuous(labels = tex_num) +
  labs(title = "Distribution of observed times", fill = NULL,
       x = "Observed time",
       y = "Count") +
  tex_theme

p <- p1 + p2
knitr::include_graphics(tikz_png(p, "figures/lab08/censoring.png",
                                 width = 10, height = 4.6))

Report.

The figure.
Whether the censored times are concentrated at the long end of the observed times, as patient B’s would be.

Exercise 2.3: Implement the EM Algorithm

Complete the E-step and M-step below.
Run EM from $\theta^{(0)} = 1$.

set.seed(536)
# EM algorithm for censored exponential
censored_exp_em <- function(Y, Delta, theta_init = 1, max_iter = 100,
                            tol = 1e-8) {
  theta <- theta_init
  n <- length(Y)
  # entry 1 holds iteration 0, the starting value
  log_lik <- c(sum(Delta) * log(theta) - theta * sum(Y),
               numeric(max_iter))
  theta_history <- c(theta, numeric(max_iter))

  for (iter in 1:max_iter) {
    theta_old <- theta

    # TODO: E-step. Compute expected survival times.
    # For uncensored (Delta==1): T_i = Y_i (we know it exactly)
    # For censored (Delta==0): E[T_i | T_i > Y_i] = Y_i + 1/theta
    E_T <- Y  # Start with observed times
    # Update E_T for censored observations (where Delta == 0)
    E_T[Delta == 0] <- ___________

    # TODO: M-step. Update theta using expected complete data.
    # theta = n / sum(E_T)
    theta <- ___________

    # Compute observed log-likelihood
    log_lik[iter + 1] <- sum(Delta) * log(theta) - theta * sum(Y)
    theta_history[iter + 1] <- theta

    # Check convergence
    if (abs(theta - theta_old) < tol) {
      log_lik <- log_lik[1:(iter + 1)]
      theta_history <- theta_history[1:(iter + 1)]
      break
    }
  }

  list(theta = theta,
       iterations = length(log_lik) - 1,
       log_lik = log_lik,
       theta_history = theta_history,
       mean_survival = 1/theta)
}

# Run EM algorithm
result_cens_em <- censored_exp_em(Y, Delta, theta_init = 1)

print("EM Algorithm Results:")
print(sprintf("Estimated theta: %.6f", result_cens_em$theta))
print(sprintf("Estimated mean survival: %.6f",
              result_cens_em$mean_survival))
print(sprintf("True theta: %.6f", true_theta))
print(sprintf("True mean survival: %.6f", 1/true_theta))
print(sprintf("Converged in %d iterations", result_cens_em$iterations))

Report.

The EM estimate of $\theta$ and of the mean survival time $1/\theta$.
The number of iterations.

Exercise 2.4: Comparison with Direct MLE

For exponential with censoring, the MLE has a closed form: \[\hat{\theta}_{MLE} = \frac{\sum_{i=1}^n \Delta_i}{\sum_{i=1}^n Y_i}\]

The code chunk below is complete. Run it as is.

set.seed(536)
# Direct MLE (closed form)
theta_mle <- sum(Delta) / sum(Y)
mean_survival_mle <- 1 / theta_mle

print("--- COMPARISON ---")
print("EM estimate:")
print(sprintf("  theta = %.6f, mean survival = %.6f",
              result_cens_em$theta, result_cens_em$mean_survival))
print("Direct MLE:")
print(sprintf("  theta = %.6f, mean survival = %.6f",
              theta_mle, mean_survival_mle))
print("True parameters:")
print(sprintf("  theta = %.6f, mean survival = %.6f",
              true_theta, 1/true_theta))
print(sprintf("Difference (EM vs MLE): %.10f",
              abs(result_cens_em$theta - theta_mle)))

Report.

The EM estimate and the direct MLE, of $\theta$ and of the mean survival time.
The difference between them.

Exercise 2.5: Visualizing Convergence

Run the figure chunk. The left panel is the observed log-likelihood at each iteration, from the starting value $\theta^{(0)} = 1$. The right panel is $\theta^{(t)}$, with the true $\theta$ and the MLE as dashed lines.

For this model the estimated standard error of the MLE is $\widehat{\text{SE}}(\hat\theta) = \hat\theta/\sqrt{\sum_i \Delta_i}$. It is one over the square root of the curvature of the log-likelihood at its peak, $-\ell''(\hat\theta) = \sum_i \Delta_i/\hat\theta^2$.

set.seed(536)
# Convergence plot
p1 <- plot_convergence(result_cens_em$log_lik,
                       "Log-likelihood convergence")

# Parameter trajectory
df_trajectory <- data.frame(
  iteration = 0:result_cens_em$iterations,
  theta = result_cens_em$theta_history,
  mean_survival = 1 / result_cens_em$theta_history
)

p2 <- ggplot(df_trajectory, aes(x = iteration, y = theta)) +
  geom_line(color = pal[["blue"]], linewidth = 1) +
  geom_point(color = pal[["red"]], size = 2) +
  geom_hline(yintercept = true_theta, linetype = "dashed",
             color = pal[["dark"]]) +
  geom_hline(yintercept = theta_mle, linetype = "dashed",
             color = pal[["green"]]) +
  annotate("text", x = result_cens_em$iterations * 0.6,
           y = true_theta + 0.04,
           label = r"(True $\theta$)", color = pal[["dark"]], size = 5.5) +
  annotate("text", x = result_cens_em$iterations * 0.6,
           y = theta_mle - 0.04,
           label = "MLE", color = pal[["green"]], size = 5.5) +
  scale_x_continuous(breaks = int_breaks, labels = tex_num) +
  scale_y_continuous(labels = tex_num) +
  labs(title = "Parameter trajectory",
       x = "Iteration",
       y = r"($\theta$)") +
  tex_theme

p <- p1 + p2
knitr::include_graphics(tikz_png(p, "figures/lab08/censored-convergence.png",
                                 width = 10, height = 4.2))

Report.

The figure.
The increase in the log-likelihood at each of the first three iterations, diff(result_cens_em$log_lik)[1:3], and the smallest increase over all iterations.
How far $\hat\theta$ is from the true $\theta = 0.5$, in estimated standard errors.

Part 3: Gaussian Mixture Model (GMM)

The model, its E-step, and its M-step are the lecture’s (handout, Section 7).

Model: Observations come from a mixture of $K$ Gaussian distributions: \[x_i \sim \sum_{k=1}^K \pi_k \Normal(\mu_k, \sigma_k^2)\]

Latent Variables: $z_i \in \{1, \ldots, K\}$ indicates which component $x_i$ came from, and $Z_{ik} = 1$ if $z_i = k$, $Z_{ik} = 0$ otherwise.

Parameters: $\theta = (\pi_1, \ldots, \pi_K, \mu_1, \ldots, \mu_K, \sigma_1^2, \ldots, \sigma_K^2)$

Complete-data log-likelihood, linear in the $Z_{ik}$: \[\ell_c(\theta; X, Z) = \sum_{i=1}^n \sum_{k=1}^K Z_{ik}\bigl[\log\pi_k + \log\phi(x_i; \mu_k, \sigma_k^2)\bigr],\] where $\phi(x; \mu, \sigma^2)$ is the Gaussian density.

EM for a Gaussian mixture

E-step: Compute the responsibilities, the posterior probabilities of the labels at $\theta^{(t)}$: \[\gamma_{ik}^{(t)} = \E_{z_i \mid x_i, \theta^{(t)}}[Z_{ik}] = P(z_i = k \mid x_i, \theta^{(t)}) = \frac{\pi_k^{(t)} \phi\bigl(x_i; \mu_k^{(t)}, (\sigma_k^{(t)})^2\bigr)}{\sum_{j=1}^K \pi_j^{(t)} \phi\bigl(x_i; \mu_j^{(t)}, (\sigma_j^{(t)})^2\bigr)}\]

M-step: With $n_k^{(t)} = \sum_{i=1}^n \gamma_{ik}^{(t)}$, update the parameters, using the responsibilities as weights: \[\pi_k^{(t+1)} = \frac{n_k^{(t)}}{n}\] \[\mu_k^{(t+1)} = \frac{\sum_{i=1}^n \gamma_{ik}^{(t)} x_i}{n_k^{(t)}}\] \[(\sigma_k^{(t+1)})^2 = \frac{\sum_{i=1}^n \gamma_{ik}^{(t)} \bigl(x_i - \mu_k^{(t+1)}\bigr)^2}{n_k^{(t)}}\]

Exercise 3.1: Generate Mixture Data

The code chunk below is complete. Run it as is.

set.seed(536)
# True parameters for K=3 mixture
K <- 3
n <- 300

true_pi <- c(0.3, 0.5, 0.2)
true_mu <- c(-2, 0, 3)
true_sigma <- c(0.5, 0.8, 0.6)

# Generate data
true_z <- sample(1:K, n, replace = TRUE, prob = true_pi)
x <- rnorm(n, mean = true_mu[true_z], sd = true_sigma[true_z])

print(sprintf("Sample size: %d, components: %d", n, K))
print(data.frame(component = 1:K, pi = true_pi, mu = true_mu,
                 sigma = true_sigma))

Exercise 3.2: Starting Values from k-means

EM needs starting values, and a common choice comes from k-means clustering. k-means splits the data into $K$ clusters by repeating two steps until no point changes cluster:

Assign each point to the nearest center.
Move each center to the mean of its points.

It is the hard-assignment version of EM for a mixture whose components have equal proportions and one common, fixed variance: $\gamma_{ik} = 1$ for the nearest center and 0 otherwise, and only the means are updated. By default kmeans() runs a refinement of these two steps (the Hartigan–Wong algorithm). nstart = 10 runs it from 10 random sets of centers and keeps the run with the smallest within-cluster sum of squares.

The two code chunks below are complete. Run them as is. The first orders the clusters by their centers and turns them into starting values: the cluster proportions, centers, and variances. The second draws the clusters next to the true components.

set.seed(536)
km <- kmeans(x, centers = K, nstart = 10)

# order the clusters by their centers, so cluster 1 has the smallest
ord <- order(km$centers[, 1])
cluster <- match(km$cluster, ord)
km_size <- as.numeric(table(cluster))

# starting values for EM
pi_init <- km_size / n
mu_init <- as.numeric(km$centers[ord, 1])
sigma2_init <- as.numeric(tapply(x, cluster, var))

print(data.frame(cluster = 1:K, size = km_size,
                 center = round(mu_init, 3), sd = round(sqrt(sigma2_init), 3)))
# true_mu is increasing, so cluster k is compared with component k
print(sprintf("Points in the cluster of their true component: %.1f%%",
              100 * mean(cluster == true_z)))

set.seed(536)
df_clusters <- data.frame(
  x = rep(x, 2),
  group = factor(c(true_z, cluster)),
  panel = factor(rep(c("True components", "k-means clusters"), each = n),
                 levels = c("True components", "k-means clusters"))
)

p <- ggplot(df_clusters, aes(x = x, fill = group)) +
  geom_histogram(bins = 40, color = "white", linewidth = 0.2) +
  facet_wrap(~ panel) +
  scale_fill_manual(values = unname(pal[c("red", "blue", "green")]),
                    name = "Component or cluster") +
  scale_x_continuous(labels = tex_num) +
  scale_y_continuous(labels = tex_num) +
  labs(x = "$x$", y = "Count") +
  tex_theme

knitr::include_graphics(tikz_png(p, "figures/lab08/mixture-kmeans.png",
                                 width = 10, height = 4.2))

Report.

The figure.
The cluster sizes and centers.
The percentage of points in the cluster of their true component.

Exercise 3.3: Implement the EM Algorithm

Complete the E-step and the M-step below.
Run EM from the k-means starting values.

set.seed(536)
# EM algorithm for a Gaussian mixture, from given starting values
gmm_em <- function(x, pi, mu, sigma2, max_iter = 1000, tol = 1e-8) {
  n <- length(x)
  K <- length(mu)

  # observed-data log-likelihood
  compute_log_lik <- function(pi, mu, sigma2) {
    dens <- sapply(1:K, function(k) pi[k] * dnorm(x, mu[k], sqrt(sigma2[k])))
    sum(log(rowSums(dens)))
  }

  # entry 1 holds iteration 0, the starting value
  log_lik <- c(compute_log_lik(pi, mu, sigma2), numeric(max_iter))

  for (iter in 1:max_iter) {
    # TODO: E-step. gamma[i, k], the responsibility of component k for x_i
    gamma <- matrix(0, n, K)
    for (k in 1:K) {
      # the numerator, for all i at once
      gamma[, k] <- ___________
    }
    # normalize, so each row sums to 1
    gamma <- ___________

    # TODO: M-step. Weighted proportions, means, and variances.
    n_k <- ___________
    pi <- ___________
    mu <- ___________
    for (k in 1:K) {
      sigma2[k] <- ___________
    }

    log_lik[iter + 1] <- compute_log_lik(pi, mu, sigma2)

    # Check convergence
    if (abs(log_lik[iter + 1] - log_lik[iter]) < tol) {
      break
    }
  }

  list(pi = pi, mu = mu, sigma2 = sigma2, gamma = gamma,
       log_lik = log_lik[1:(iter + 1)], iterations = iter)
}

# Run EM from the k-means starting values
result_gmm <- gmm_em(x, pi_init, mu_init, sigma2_init)

print(sprintf("Converged in %d iterations", result_gmm$iterations))
print(sprintf("Final log-likelihood: %.4f", tail(result_gmm$log_lik, 1)))
print(data.frame(component = 1:K,
                 pi = round(result_gmm$pi, 3), true_pi = true_pi,
                 mu = round(result_gmm$mu, 3), true_mu = true_mu,
                 sigma = round(sqrt(result_gmm$sigma2), 3),
                 true_sigma = true_sigma))

Report.

The estimates, next to the true values.
The number of iterations and the final log-likelihood.

Exercise 3.4: Identical Starting Values

Run EM with every component started at the same values: $\pi_k^{(0)} = 1/3$, $\mu_k^{(0)} = \bar x$, and $(\sigma_k^{(0)})^2 = s^2$, the sample variance.
Run it again with the three means nudged apart by 0.01.

The code chunk below is complete. Run it as is.

set.seed(536)
# every component starts at the same mean and variance
result_same <- gmm_em(x, pi = rep(1/K, K), mu = rep(mean(x), K),
                      sigma2 = rep(var(x), K))
# the same start, with the means nudged apart by 0.01
result_nudged <- gmm_em(x, pi = rep(1/K, K), mu = mean(x) + c(-0.01, 0, 0.01),
                        sigma2 = rep(var(x), K))
# the MLE of one normal, for comparison
ll_single <- sum(dnorm(x, mean(x), sqrt(mean((x - mean(x))^2)), log = TRUE))

print(sprintf("Identical start: %d iterations, log-likelihood %.4f",
              result_same$iterations, tail(result_same$log_lik, 1)))
print(data.frame(component = 1:K, pi = result_same$pi, mu = result_same$mu,
                 sigma = sqrt(result_same$sigma2)))
print(sprintf("One normal (MLE): log-likelihood %.4f", ll_single))
print(sprintf("Nudged start: %d iterations, log-likelihood %.4f",
              result_nudged$iterations, tail(result_nudged$log_lik, 1)))
print(sprintf("k-means start: %d iterations, log-likelihood %.4f",
              result_gmm$iterations, tail(result_gmm$log_lik, 1)))

Report.

What EM returns from the identical start, and after how many iterations.
Its final log-likelihood, next to that of one normal and the final one of the run from the k-means start.
What the nudge changes.

Exercise 3.5: Visualizing the Fits

Run the figure chunk. The left panel is the log-likelihood at each iteration, from the three starts. The right panel is the histogram of the data, with the true mixture density and two fitted ones.

set.seed(536)
# Log-likelihood paths
paths <- list("k-means start" = result_gmm$log_lik,
              "identical start" = result_same$log_lik,
              "nudged start" = result_nudged$log_lik)
df_paths <- do.call(rbind, lapply(names(paths), function(s)
  data.frame(iteration = seq_along(paths[[s]]) - 1, log_lik = paths[[s]],
             start = s)))
df_paths$start <- factor(df_paths$start, levels = names(paths))

p1 <- ggplot(df_paths, aes(x = iteration, y = log_lik, color = start)) +
  geom_line(linewidth = 1) +
  geom_point(data = subset(df_paths, start == "identical start"), size = 2.5) +
  scale_color_manual(values = unname(pal[c("blue", "red", "orange")]),
                     name = NULL) +
  scale_x_continuous(breaks = int_breaks, labels = tex_num) +
  scale_y_continuous(labels = tex_num) +
  labs(title = "Log-likelihood by start", x = "Iteration",
       y = "Log-likelihood") +
  tex_theme +
  theme(legend.position = "inside", legend.position.inside = c(0.68, 0.45))

# Fitted densities
mixture_density <- function(x0, pi, mu, sigma) {
  rowSums(sapply(seq_along(mu), function(k) pi[k] * dnorm(x0, mu[k], sigma[k])))
}
x_seq <- seq(min(x) - 1, max(x) + 1, length.out = 400)
fits <- c("true mixture", "k-means start", "identical start")
df_dens <- data.frame(
  x = rep(x_seq, 3),
  density = c(mixture_density(x_seq, true_pi, true_mu, true_sigma),
              mixture_density(x_seq, result_gmm$pi, result_gmm$mu,
                              sqrt(result_gmm$sigma2)),
              mixture_density(x_seq, result_same$pi, result_same$mu,
                              sqrt(result_same$sigma2))),
  fit = factor(rep(fits, each = length(x_seq)), levels = fits)
)

p2 <- ggplot() +
  geom_histogram(data = data.frame(x = x), aes(x = x, y = after_stat(density)),
                 bins = 40, fill = "grey85", color = "white") +
  geom_line(data = df_dens, aes(x = x, y = density, color = fit, linetype = fit),
            linewidth = 1) +
  scale_color_manual(values = unname(pal[c("dark", "blue", "red")]),
                     name = NULL) +
  scale_linetype_manual(values = c("dashed", "solid", "solid"), name = NULL) +
  scale_x_continuous(labels = tex_num) +
  scale_y_continuous(labels = tex_num) +
  labs(title = "Fitted densities", x = "$x$", y = "Density") +
  tex_theme +
  theme(legend.position = "inside", legend.position.inside = c(0.8, 0.84))

p <- p1 + p2
knitr::include_graphics(tikz_png(p, "figures/lab08/mixture-fits.png",
                                 width = 11, height = 4.6))

Report.

The figure.
The first iteration at which the nudged path is more than 1 above the identical start’s log-likelihood, min(which(result_nudged$log_lik > ll_single + 1)) - 1.
The largest gap between each fitted density and the true one, over the plotted range: for the k-means start, with(df_dens, max(abs(density[fit == "k-means start"] - density[fit == "true mixture"]))), and the same with "identical start".

Summary and Reflection

Takeaways

In all three problems, the E-step reduces to expected complete data. In the linkage problem it splits the mixed cell: $\E_{Z_2 \mid X_1, \psi^{(t)}}[Z_2]$. For censored data it replaces each censored time by $Y_i + 1/\theta^{(t)}$, from the memoryless property. In the mixture it replaces each label indicator by its responsibility $\gamma_{ik}^{(t)}$.
The M-step uses these expectations as if they were real. In all three problems the terms of $\ell_c$ that involve the parameter are linear in the missing data, so the complete-data MLE with the expectations plugged in is the maximizer of $Q$.
The log-likelihood never decreases. The monotonicity theorem guarantees it at every iteration, from the starting value on.
In Parts 1 and 2, EM converges to the MLE. Both log-likelihoods are strictly concave, so their only stationary point is the MLE. In general, EM’s limit is only a stationary point. A closed-form MLE or a direct optimizer, where one exists, is a check on the EM code.
EM and BFGS count their work differently. Here an EM iteration is a closed-form update. A BFGS iteration needs a gradient, here by finite differences, and a line search.
The MLE carries sampling error. EM finds the MLE here. The distance between the MLE and the true parameter is measured in estimated standard errors.
For a mixture, the start decides where EM ends. Its log-likelihood is not concave and has stationary points that are not maxima. Start from k-means or from several random starts, and compare the final log-likelihoods.

Self-Assessment Questions

Explain the difference between the E-step and the M-step.
Explain why EM monotonically increases the likelihood.
True/False: EM always converges to the global maximum. Explain.
For Part 2, write $Q(\theta \mid \theta^{(t)})$ from $\ell_c(\theta; T) = n\log\theta - \theta\sum_i T_i$ and maximize it. Check that you get the M-step of the code.
optim() reports counts of function and gradient evaluations. Which one compares with the number of EM iterations, and why?
Describe when EM is preferable to direct optimization.
Show from the E-step that if every component starts with the same mean and variance, then $\gamma_{ik}^{(t)} = \pi_k^{(t)}$, and the M-step gives every component the same mean and variance again.
What does k-means compute, and how is it related to EM for a Gaussian mixture?

Symbol	Meaning
\(X\); \(Z\)	the observed data; the latent (missing) data
\(\theta\)	the parameter (\(\psi\) in Part 1)
\(\theta^{(t)}\)	the parameter after \(t\) iterations; \(\theta^{(0)}\) is the starting value
\(\ell(\theta; X)\)	the observed-data log-likelihood
\(\ell_c(\theta; X, Z) = \log p(X, Z \mid \theta)\)	the complete-data log-likelihood
\(p(Z \mid X, \theta^{(t)})\)	the posterior of \(Z\) at \(\theta^{(t)}\)
\(Q(\theta \mid \theta^{(t)})\)	the expected complete-data log-likelihood