ECE 60141: Foundations of Computational Imaging

Lab 6: The EM Algorithm

Overview What the lab is about, and the ground rules.

What This Lab Is About

Due: Friday, Nov. 13.

Every estimate you have computed in this course used data you could see. The EM algorithm computes maximum likelihood estimates when part of the data is missing, and here the missing part is the answer key. Each data point came from one of several Gaussian clusters, and nobody records which. EM estimates the cluster parameters anyway, which makes it a clustering algorithm. Then you will use it to answer a harder question that has no obvious likelihood at all: how many clusters are there?

This lab also teaches something about building software with an AI. In the earlier labs the page handed you the exact signature of every function. This one does not, and that is deliberate. Here the program has real structure, so you will design it yourself: for each piece you write the specification, saying what it computes, what goes in, and what comes out, and only then does the AI write the code. Writing a specification precise enough that a stranger could implement it is the skill this lab grades.

The background is Chapter 12. The notation here matches the book.

Ground Rules

  • Use the labs conda environment from Lab 2. All code in python, using numpy or PyTorch.
  • Every experiment on this page is fully specified: the data, the initialization, and the number of iterations. Follow it exactly, so your numbers can be compared with everyone else's.
  • Seed every random number generator with 0 unless the step says otherwise, for example rng = np.random.default_rng(0). State in your report which generator you used.
  • Each function comes with checks. Run them. If a check fails, the code is wrong, and it is your job to get it fixed.
  • Keep your code in your image-processing-labs repository, in a directory named lab6.
Directing an AI well. Hand it one specification at a time, together with the equations from this page that it implements. Then read the code and compare it line by line against those equations. Code that runs without an error message is not the same as code that is correct, and this lab has two failure modes that produce no error message at all. Step 7 names them.
Module 1: The Model and the Algorithm Steps 1 to 3. Where EM comes from, in three short steps.

Step 1: The Model

Let \( \{X_n\}_{n=1}^{N} \) be i.i.d. discrete labels taking values in \( \{0, \cdots, M-1\} \) with \( P\{X_n = m\} = \pi_m \). Let \( \{Y_n\}_{n=1}^{N} \), with \( Y_n \in \mathbb{R}^p \), be conditionally independent Gaussian random vectors given the labels, with \( Y_n \sim N(\mu_{x_n}, R_{x_n}) \). Write \( \theta_m = (\mu_m, R_m) \) for the parameters of component \( m \), so that

$$ p(y_n | \theta_m) = \frac{1}{(2\pi)^{p/2} |R_m|^{1/2}} \exp\left\{ -\frac{1}{2} (y_n - \mu_m)^t R_m^{-1} (y_n - \mu_m) \right\} , \tag{1} $$

and the full parameter vector is \( \theta = (\pi_0, \mu_0, R_0, \cdots, \pi_{M-1}, \mu_{M-1}, R_{M-1}) \).

The labels are never observed. Summing them out gives the Gaussian mixture distribution

$$ p(y_n | \theta) = \sum_{m=0}^{M-1} \pi_m \, p(y_n | \theta_m) , \tag{2} $$

and, since the samples are i.i.d., the log likelihood of the whole data set is

$$ \log p(y | M, \theta) = \sum_{n=1}^{N} \log \left( \sum_{m=0}^{M-1} \pi_m \, p(y_n | \theta_m) \right) . \tag{3} $$

We want the \( \theta \) that maximizes (3). Direct maximization is hard, and the reason is visible in the formula: there is a sum inside the logarithm, so the log does not distribute and nothing separates.

Step 2: What You Could Do If You Knew the Labels

Suppose for a moment that the labels \( x_n \) were observed. Then estimation would be easy, and it is worth seeing how easy, because EM is built directly on this case.

Everything the ML estimate needs to know about labeled data is contained in three statistics per class:

$$ N_m = \sum_{n=1}^{N} \delta(x_n - m) , \qquad b_m = \sum_{n=1}^{N} y_n \, \delta(x_n - m) , \qquad S_m = \sum_{n=1}^{N} y_n y_n^t \, \delta(x_n - m) , \tag{4} $$

where \( \delta(\cdot) \) is 1 at zero and 0 elsewhere. So \( N_m \) counts the points in class \( m \), \( b_m \) sums them, and \( S_m \) sums their outer products. Chapter 12 shows that these are the natural sufficient statistics of the complete-data distribution, and that the ML estimates are simple functions of them:

$$ \hat{\pi}_m = \frac{N_m}{N} , \qquad \hat{\mu}_m = \frac{b_m}{N_m} , \qquad \hat{R}_m = \frac{S_m}{N_m} - \frac{b_m b_m^t}{N_m^2} . \tag{5} $$

These are just the class fraction, the class sample mean, and the class sample covariance. Keep (4) and (5) in view for the next step.

Step 3: The EM Updates

The labels are not observed, so the statistics (4) cannot be computed. The EM idea is one sentence: replace each statistic by its expected value given the observed data and the current parameter estimate, then apply (5) as if the expected statistics were the real thing.

The expectation of \( \delta(X_n - m) \) given \( Y = y \) is the posterior probability that sample \( n \) belongs to class \( m \). The E-step computes it for every \( n \) and \( m \):

$$ P_{n,m} \leftarrow P\{ X_n = m \,|\, Y = y, \theta \} = \frac{ \pi_m \, p(y_n | \theta_m) } { \displaystyle \sum_{j=0}^{M-1} \pi_j \, p(y_n | \theta_j) } . \tag{6} $$

These are the soft assignments. Sample \( n \) has partial membership \( P_{n,m} \) in cluster \( m \), and \( \sum_m P_{n,m} = 1 \). The expected statistics are then (4) with \( \delta(x_n - m) \) replaced by \( P_{n,m} \):

$$ \bar{N}_m \leftarrow \sum_{n=1}^{N} P_{n,m} , \qquad \bar{b}_m \leftarrow \sum_{n=1}^{N} y_n P_{n,m} , \qquad \bar{S}_m \leftarrow \sum_{n=1}^{N} y_n y_n^t P_{n,m} , \tag{7} $$

and the M-step is (5) applied to the expected statistics:

$$ \pi_m \leftarrow \frac{\bar{N}_m}{N} , \qquad \mu_m \leftarrow \frac{\bar{b}_m}{\bar{N}_m} , \qquad R_m \leftarrow \frac{\bar{S}_m}{\bar{N}_m} - \frac{\bar{b}_m \bar{b}_m^t}{\bar{N}_m^2} . \tag{8} $$

One EM iteration is one E-step followed by one M-step. Chapter 12 proves the property that makes the algorithm usable: each iteration increases the log likelihood (3), or leaves it unchanged. It never decreases. That property is also the best test you have of whether your code is right, and you will check it numerically in Step 7.

Module 2: Specify It, Then Build It Steps 4 to 7. The data, two derivations by hand, and four specifications.

Step 4: Generate the Data

Generate \( N = 500 \) samples from a Gaussian mixture with \( p = 2 \), \( M = 3 \), and \( \pi = [0.4, 0.4, 0.2] \), means

$$ \mu_0 = \begin{bmatrix} 2 \\ 2 \end{bmatrix} , \qquad \mu_1 = \begin{bmatrix} -2 \\ -2 \end{bmatrix} , \qquad \mu_2 = \begin{bmatrix} 5.5 \\ 2 \end{bmatrix} , \tag{9} $$

and covariances

$$ R_0 = \begin{bmatrix} 1 & 0.1 \\ 0.1 & 1 \end{bmatrix} , \qquad R_1 = \begin{bmatrix} 1 & -0.1 \\ -0.1 & 1 \end{bmatrix} , \qquad R_2 = \begin{bmatrix} 1 & 0.2 \\ 0.2 & 0.5 \end{bmatrix} . \tag{10} $$

Draw each label \( x_n \) from \( \pi \), then draw \( y_n \) from \( N(\mu_{x_n}, R_{x_n}) \). Save both \( y \), of shape \( 500 \times 2 \), and the true labels \( x \) to a file. You will reuse this exact data set in every remaining step, so generate it once and load it after that.

HAND IN · D1

A scatter plot of the 500 samples, with the three true classes in three colors and a legend. Use equal axis scaling, so that the covariances are not distorted by the aspect ratio.

Step 5: Derive Two Things by Hand

Do this step without AI help. Both derivations are short, and they are what make the rest of the lab make sense rather than being a sequence of formulas to type in.

  1. The E-step. Derive equation (6). Start from Bayes' rule for \( P\{X_n = m \,|\, Y = y, \theta\} \), and use the fact that the pairs \( (Y_n, X_n) \) are independent across \( n \) to show that the posterior for \( X_n \) depends on \( y \) only through \( y_n \).
  2. The M-step reduces to Step 2. Show that if the soft assignments are hard, meaning \( P_{n,m} = \delta(x_n - m) \) for known labels \( x_n \), then the expected statistics (7) equal the complete-data statistics (4), so the M-step (8) reproduces the labeled ML estimates (5) exactly.
The second derivation is also a test. You will turn it into a numerical check in Step 7. That is the general habit worth taking from this lab: a derivation that ends in an equality is a test case waiting to be written.
HAND IN · D2

Both derivations, handwritten or typeset, a few lines each.

Step 6: Write the Specifications

Now design the implementation. Do not ask the AI for "an EM program." Break the algorithm into the four functions below, and write a specification for each one before any code exists. A specification has three parts:

  • What it does, in a few sentences, naming the equation numbers from this page that it implements.
  • Inputs: every input, with its meaning, type, and array shape.
  • Outputs: every output, with its meaning, type, and array shape.

The decomposition is fixed, so that everyone's code has the same structure. The interfaces, meaning the exact inputs, outputs, and shapes, are yours to design.

FunctionWhat it computes
e_step The posterior class probabilities \( P_{n,m} \) of equation (6), for every sample and every class.
m_step The expected statistics \( \bar{N}_m \), \( \bar{b}_m \), \( \bar{S}_m \) of equation (7), and from them the updated parameters of equation (8).
log_likelihood The log likelihood of equation (3), for a given data set and parameter vector.
run_em The outer loop. Alternates e_step and m_step until the log likelihood stops improving by more than a tolerance, or a maximum iteration count is reached. Records the log likelihood after every iteration.
A test of a specification. Could a classmate who has read this page, but cannot ask you any questions, implement your function and have it work with the other three? If not, the specification is incomplete. The usual thing missing is a shape.
HAND IN · D3

The four specifications, verbatim. These are a graded part of the lab, not a preliminary.

Step 7: Build and Test

Hand the AI one specification at a time, and test each function before you build the next.

Two implementation requirements apply to everyone, and both exist because of failures that produce no error message.

  • Compute (3) and the E-step (6) in the log domain, using a log-sum-exp function. Computing \( p(y_n|\theta_m) \) directly underflows to zero for points far from a component mean, and then (6) divides zero by zero.
  • After each covariance update in (8), add \( \epsilon I \) to \( R_m \) with \( \epsilon = 10^{-6} \). Without it, a component that collects only a few points produces a singular covariance and the log likelihood runs off to infinity.

Run these four checks.

  1. e_step. Every row of \( P_{n,m} \) sums to one, to within numerical precision. Include a test point placed far from all component means, where an implementation that skipped the log domain returns NaN.
  2. m_step. Feed it the hard assignments \( P_{n,m} = \delta(x_n - m) \) built from the true labels of Step 4. By your second derivation in Step 5, the output must equal the class fractions, class sample means, and class sample covariances computed directly from the labeled data. Verify that it does.
  3. log_likelihood. For \( M = 1 \), the mixture (2) is a single Gaussian. Check the function against a direct evaluation of the Gaussian log density summed over the samples.
  4. run_em. The recorded log likelihood is nondecreasing from one iteration to the next. Report the largest decrease you observe, which should be zero or a rounding-level negative number. Also check that if you initialize at the true parameters of Step 4, the parameters stay close to the true values.
HAND IN · D4

The result of each of the four checks, with the numbers.

Module 3: Run It, and Watch It Get Stuck Steps 8 and 9. Clustering, and what happens when you start somewhere else.

Step 8: Run EM

Run EM on the data from Step 4, with \( M = 3 \), exactly 20 iterations with the tolerance set to zero, and the initialization

$$ \pi \leftarrow [1/3, 1/3, 1/3] , \qquad R_m \leftarrow \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix} \ \mbox{ for } m = 0, 1, 2 , \tag{11} $$

with \( \mu_0, \mu_1, \mu_2 \) set to the first three sample vectors produced in Step 4.

The labels EM produces are not the labels you generated. The mixture likelihood does not care which component is called 0, so EM will land on some permutation of the components. Before you compare anything, find the correspondence between estimated and true components that best matches their means, and apply it.
HAND IN · D5

A plot of \( \log p(y | M, \theta) \) versus iteration number, for iterations 1 through 20.

HAND IN · D6

A table of the estimated \( \pi_m \), \( \mu_m \), and \( R_m \), each placed next to the true value it corresponds to, using the correspondence you found.

HAND IN · D7

A scatter plot of the data colored by the hard label \( \hat{x}_n = \arg\max_m P_{n,m} \), with the estimated mean of each component marked, and the fraction of samples whose hard label disagrees with the true label from Step 4.

Step 9: Initialization and Local Maxima

The log likelihood of a mixture has many local maxima, and EM finds only the one nearest to where it starts. Every iteration goes uphill, the algorithm is working exactly as proved, and the answer can still be poor. That is worth seeing once, deliberately, rather than discovering it in a project.

Repeat Step 8 five more times with \( M = 3 \), keeping \( \pi \) and \( R_m \) as in (11), but choosing the three initial means as three distinct data points drawn at random from the 500 samples. Use seeds 0, 1, 2, 3, 4 for the five runs. Run 20 iterations each.

HAND IN · D8

One plot with all five log likelihood curves plus the curve from Step 8, on the same axes with a legend, and a table of the final log likelihood for all six runs.

HAND IN · D9

The hard-label scatter plot for whichever run reached the lowest final log likelihood, with one sentence describing what that solution did with the clusters.

Module 4: How Many Clusters Are There? Steps 10 and 11. A question the likelihood cannot answer on its own.

Step 10: Choosing the Number of Clusters

So far you told the algorithm that \( M = 3 \). In practice nobody tells you. And you cannot choose \( M \) by maximizing the likelihood, because any order-\( M \) model is a special case of an order-\( (M+1) \) model, so the maximized likelihood can only go up with \( M \). Left alone, the likelihood will happily give every data point its own cluster.

The fix is to add a penalty that grows with the number of parameters. The minimum description length criterion is

$$ MDL(M, \theta) = -\log p(y | M, \theta) + \frac{1}{2} L(M) \log(Np) , \tag{12} $$

where \( L(M) \) is the number of real numbers needed to specify \( \theta \). Each component needs 1 number for \( \pi_m \), \( p \) numbers for \( \mu_m \), and \( (p+1)p/2 \) numbers for the symmetric matrix \( R_m \), and one number is redundant because the \( \pi_m \) sum to one. So

$$ L(M) = M \left( 1 + p + \frac{(p+1)p}{2} \right) - 1 . \tag{13} $$

To find \( \hat{M} \), start with more components than you expect and remove them one at a time.

Order estimation by merging

  Initialize M_o to the largest order you will consider, and
  initialize theta
  For M = M_o down to 1:
      Run EM to produce theta*(M)
      Record MDL(M) using (12)
      Merge two components of theta*(M) to initialize order M-1
  Select M_hat as the value of M minimizing MDL(M)

Merge the pair of components \( (l, m) \) that minimizes the distance

$$ d(l,m) = \frac{N \pi_l}{2} \log\left( \frac{|R_{(l,m)}|}{|R_l|} \right) + \frac{N \pi_m}{2} \log\left( \frac{|R_{(l,m)}|}{|R_m|} \right) , \tag{14} $$

where the merged component has parameters

$$ \pi_{(l,m)} = \pi_l + \pi_m , \qquad \mu_{(l,m)} = \frac{\pi_l \mu_l + \pi_m \mu_m}{\pi_l + \pi_m} , \tag{15} $$ $$ R_{(l,m)} = \frac{ \pi_l \left( R_l + (\mu_l - \mu_{(l,m)})(\mu_l - \mu_{(l,m)})^t \right) + \pi_m \left( R_m + (\mu_m - \mu_{(l,m)})(\mu_m - \mu_{(l,m)})^t \right) }{\pi_l + \pi_m} . \tag{16} $$

This distance upper bounds the increase in MDL caused by the merge, so merging the closest pair gives up as little as possible at each reduction in order. After the merge, the merged parameters become the new component and the others carry over unchanged.

Three more functions are needed, and you specify them the same way you did in Step 6, saying what they compute, what goes in, and what comes out, before any code.

FunctionWhat it computes
mdl The MDL value of equations (12) and (13), using your log_likelihood.
merge_closest Finds the pair minimizing (14), merges it using (15) and (16), and returns the order-\( (M-1) \) parameter vector.
estimate_order The search loop above, using run_em, mdl, and merge_closest.

Run it on the data from Step 4 with \( M_o = 9 \), the initialization \( \pi \leftarrow [1/9, \cdots, 1/9] \), \( R_m \leftarrow I \) for all \( m \), and \( \mu_0 \) through \( \mu_8 \) set to the first nine sample vectors from Step 4. Run 20 EM iterations for each value of \( M \).

HAND IN · D10

The three additional specifications, verbatim.

HAND IN · D11

A plot of the MDL value versus the cumulative number of EM iterations over the whole run. Plot the MDL value after each EM iteration, and when an iteration is followed by a merge, also plot the value immediately after the merge, so that merge steps show two points. Mark the merge locations.

HAND IN · D12

A plot of \( MDL(M) \) versus \( M \) for \( M = 1, \cdots, 9 \), with the minimizing \( \hat{M} \) labeled, and a statement of whether \( \hat{M} \) equals the true value \( M = 3 \).

Step 11: Questions to Answer

Answer each in a few sentences, using your own plots and numbers as evidence.

  1. Compare the estimated parameters from Step 8 to the true values. Which are estimated most accurately, and which least? Give a reason based on the data, not on the algorithm.
  2. In Step 9, did every initialization reach the same final log likelihood? What does the spread of final values tell you about the shape of the likelihood surface, and what would you do about it in a real application?
  3. The likelihood of a Gaussian mixture has no finite maximum. You can drive it to infinity by putting a component on a single data point and shrinking its covariance. Explain how the \( \epsilon I \) term of Step 7 prevents this, and what you give up by using it.
  4. In (12) the penalty grows like \( \log(Np) \). What happens to \( \hat{M} \) as you collect more data from the same distribution, and why is that the behavior you want?
  5. Step 10 estimates \( \theta \) for each \( M \) by starting from a merge of the previous solution rather than from a fresh initialization. Give one advantage and one risk of doing it that way.
  6. Which of your seven specifications needed revision after you saw the code or a failed check? What was missing from the original version?
HAND IN · D13

Your answers to the six questions above.

Deliverables What to submit, and where each item comes from.

What to Hand In

Submit two files through Brightspace.

  1. A report, as a single PDF, labeled "Lab 6", containing the items below in this order.
  2. Your code, as plain .py files. Do not paste the code into the report and do not send a zip file.

Also commit your code to your image-processing-labs repository under lab6 and push it before you submit.

  1. Your name, the link to your GitHub repository, and which random number generator you used.
  2. D1 The scatter plot of the generated data.
  3. D2 Your two hand derivations.
  4. D3 Your four specifications, verbatim.
  5. D4 The results of the four checks, with the numbers.
  6. D5 The log likelihood plot for the 20 iterations.
  7. D6 The estimated against true parameter table.
  8. D7 The hard-label scatter plot, with the disagreement fraction.
  9. D8 The six log likelihood curves and the table of final values.
  10. D9 The scatter plot for the worst run, with one sentence.
  11. D10 The three additional specifications, verbatim.
  12. D11 The MDL against cumulative EM iterations plot, with the merges marked.
  13. D12 The MDL against \( M \) plot, with \( \hat{M} \) labeled and the comparison to the true value.
  14. D13 Your answers to the six questions.

Concise and clear beats long.

EM is the standard tool for fitting models with hidden variables, and the pattern you used here, expected sufficient statistics in place of observed ones, carries far beyond Gaussian mixtures. So does the other habit this lab asked for: writing down what a piece of code must do before anything writes it.

Back to the laboratory index · ECE 60141 course page