Expectation Maximization

Two steps in a loop

Suppose a box holds buttons from two jars, but nobody labels which jar a button came from. You can still learn both jars. Guess which jar each button belongs to, then move each jar to the average of the buttons assigned to it. Repeat until nothing changes. That is the idea behind expectation maximization (EM), a classic training method for models with latent variables , the hidden labels.

EM alternates two steps. The E-step sets the approximate posterior from the current parameters. The M-step changes the parameters to raise the bound with held fixed. The book describes this as a view developed by Neal and Hinton, and it has one useful property: EM is not an approximate inference method in itself. It is a way of learning with an approximate posterior.

Here is a toy version. Forty values were drawn from two bell curves, one centred near −2 and one near +2, using numpy’s default_rng with seed 7, rounded to two decimals. The program does not know which curve each value came from. Its two guesses are the means (teal) and (violet). Start EM with the buttons below.

The dot sizes in the lower lanes are the responsibilities , the share of each data value that the E-step gives to each cluster. Each value in the lower lanes has two dots, one per cluster, and the two shares add up to 1. Values near a mean get large responsibilities for that mean. Values in between are shared.

The E-step closes the gap

The E-step fixes the approximate posterior at the current parameters, which the book writes as for the value at the start of the step. Once is fixed, the bound touches the true log probability at that . Moving away from keeps unchanged, so the bound falls behind the true value. The gap then grows.

The M-step assumes the same works for every new . That assumption opens a gap between the bound and the true log probability. The next E-step closes it again, because it sets to the posterior at the new parameters. The E-step itself is exact inference, yet the book notes that EM can still be read as using approximate inference in some sense, because of this reuse of across parameter values.

The gap is the difference between the true log probability and the bound, summed over all examples. This sum is added here and is not printed in the book.

The M-step and a partial step

The M-step chooses to maximize the sum of the bound over the training examples, with held fixed:

The book says that this maximization may be complete or partial. A partial M-step only moves part of the way towards the best value. Either choice still raises the bound, and EM only needs that.

For this toy, the maximization has a closed form. The weights are the responsibilities from the E-step, and each mean becomes a responsibility-weighted average of the data. This derivation is added here and is not printed in the book. It takes the derivative of the sum with respect to and sets it to zero.

The demo starts from the means (−0.5, 0.5), keeps the responsibilities from that E-step, and moves the means a fraction of the way to the closed-form answer. The bound rises at every fraction, but the true log-likelihood rises faster over the whole step, so the gap grows with the fraction. That is why the gap opens after an M-step.

EM as coordinate ascent

EM can be read as coordinate ascent on the bound . The E-step maximizes it over the approximate posterior , and the M-step maximizes it over the parameters . Each step lifts the bound, so the process settles at a point, but that point is not always the best possible one.

The book names two insights in EM. The first is the loop itself: update the parameters to raise the likelihood of a completed data set, in which each missing label is filled in by an estimate of its posterior. Gradient ascent on the log-likelihood has the same property, because its gradient is an expectation under the posterior. The second is that one can be reused after has moved, which is what allows large M-steps.

The book adds three points about the M-step. Stochastic gradient ascent on latent variable models is a special case of EM in which the M-step is a single gradient step. Other forms can take much larger M-steps, and for some families the M-step has an exact answer. In deep learning, though, most models are too complicated for an exact large M-step, so the second insight is rarely used. Test yourself on the four actions below.

Check yourself

Questions that mix the whole module. Each one is scored, and the score feeds your mastery of this module.

Source: Ian Goodfellow, Yoshua Bengio and Aaron Courville, Deep Learning, MIT Press, 2016, Section 19.2 (pages 634–635), free to read at deeplearningbook.org. Where the book prints equation numbers, the numbers here follow the book. The order of ideas follows the book. The explanations and all demos were written for this page.