In this set of notes, we discuss the EM (Expectation-Maximization) algorithm for density estimation.

Suppose that we are given a training set $\{x^{(1)},\dots ,x^{(n)}\}$$\left\{x^{(1)},\dots ,x^{(n)}\right\}${x^((1)),dots,x^((n))}\left\{x^{(1)}, \ldots, x^{(n)}\right\} as usual. Since we are in the unsupervised learning setting, these points do not come with any labels.

We wish to model the data by specifying a joint distribution $p(x^{(i)},z^{(i)})=$$p\left(x^{(i)},z^{(i)}\right)=$p(x^((i)),z^((i)))=p\left(x^{(i)}, z^{(i)}\right)= $p\left(x^{(i)}|z^{(i)}\right)p\left(z^{(i)}\right)$$p\left(x^{(i)}|z^{(i)}\right)p\left(z^{(i)}\right)$p(x^((i))|z^((i)))p(z^((i)))p\left(x^{(i)} | z^{(i)}\right) p\left(z^{(i)}\right). Here, $z^{(i)}\sim \mathrm{Multinomial}(\varphi )$$z^{(i)}\sim \mathrm{Multinomial}(\varphi )$z^((i))∼Multinomial(phi)z^{(i)} \sim \operatorname{Multinomial}(\phi) (where ${\varphi }_j\ge 0,\sum _{j=1}^k{\varphi }_j=1$${\varphi }_j\ge 0,\sum _{j=1}^k {\varphi }_j=1$phi_(j) >= 0,sum_(j=1)^(k)phi_(j)=1\phi_{j} \geq 0, \sum_{j=1}^{k} \phi_{j}=1, and the parameter ${\varphi }_j$${\varphi }_j$phi_(j)\phi_{j} gives $p(z^{(i)}=j))$$\left.p\left(z^{(i)}=j\right)\right)${:p(z^((i))=j))\left.p\left(z^{(i)}=j\right)\right), and $x^{(i)}|z^{(i)}=j\sim \mathcal{N}({\mu }_j,{\mathrm{\Sigma }}_j)$$x^{(i)}|z^{(i)}=j\sim \mathcal{N}\left({\mu }_j,{\mathrm{\Sigma }}_j\right)$x^((i))|z^((i))=j∼N(mu_(j),Sigma_(j))x^{(i)} | z^{(i)}=j \sim \mathcal{N}\left(\mu_{j}, \Sigma_{j}\right). We let $k$$k$kk denote the number of values that the $z^{(i)}$$z^{(i)}$z^((i))z^{(i)}'s can take on. Thus, our model posits that each $x^{(i)}$$x^{(i)}$x^((i))x^{(i)} was generated by randomly choosing $z^{(i)}$$z^{(i)}$z^((i))z^{(i)} from $\{1,\dots ,k\}$$\{1,\dots ,k\}${1,dots,k}\{1, \ldots, k\}, and then $x^{(i)}$$x^{(i)}$x^((i))x^{(i)} was drawn from one of $k$$k$kk Gaussians depending on $z^{(i)}$$z^{(i)}$z^((i))z^{(i)}. This is called the mixture of Gaussians model. Also, note that the $z^{(i)}$$z^{(i)}$z^((i))z^{(i)}'s are latent random variables, meaning that they're hidden/unobserved. This is what will make our estimation problem difficult.

The parameters of our model are thus $\varphi ,\mu$$\varphi ,\mu$phi,mu\phi, \mu and $\mathrm{\Sigma }$$\mathrm{\Sigma }$Sigma\Sigma. To estimate them, we can write down the likelihood of our data:

However, if we set to zero the derivatives of this formula with respect to the parameters and try to solve, we'll find that it is not possible to find the maximum likelihood estimates of the parameters in closed form. (Try this yourself at home.)

The random variables $z^{(i)}$$z^{(i)}$z^((i))z^{(i)} indicate which of the $k$$k$kk Gaussians each $x^{(i)}$$x^{(i)}$x^((i))x^{(i)} had come from. Note that if we knew what the $z^{(i)}$$z^{(i)}$z^((i))z^{(i)}'s were, the maximum likelihood problem would have been easy. Specifically, we could then write down the likelihood as

Maximizing this with respect to $\varphi ,\mu$$\varphi ,\mu$phi,mu\phi, \mu and $\mathrm{\Sigma }$$\mathrm{\Sigma }$Sigma\Sigma gives the parameters:

Indeed, we see that if the $z^{(i)}$$z^{(i)}$z^((i))z^{(i)}'s were known, then maximum likelihood estimation becomes nearly identical to what we had when estimating the parameters of the Gaussian discriminant analysis model, except that here the $z^{(i)}$$z^{(i)}$z^((i))z^{(i)}'s playing the role of the class labels.^[1]

However, in our density estimation problem, the $z^{(i)}$$z^{(i)}$z^((i))z^{(i)}’s are not known. What can we do?

The EM algorithm is an iterative algorithm that has two main steps. Applied to our problem, in the E-step, it tries to "guess" the values of the $z^{(i)}$$z^{(i)}$z^((i))z^{(i)}'s. In the M-step, it updates the parameters of our model based on our guesses. Since in the M-step we are pretending that the guesses in the first part were correct, the maximization becomes easy. Here's the algorithm:

$$$$quad\quadRepeat until convergence: $\{$$\{${\{

$$$$quadquad\quad\quad(E-step) For each $i,j,$$i,j,$i,j,i, j, set

$$$$quadquad\quad\quad(M-step) Update the parameters:

In the E-step, we calculate the posterior probability of our parameters the $z^{(i)}$$z^{(i)}$z^((i))z^{(i)}'s, given the $x^{(i)}$$x^{(i)}$x^((i))x^{(i)} and using the current setting of our parameters. I.e., using Bayes rule, we obtain:

Here, $p(x^{(i)}|z^{(i)}=j;\mu ,\mathrm{\Sigma })$$p\left(x^{(i)}|z^{(i)}=j;\mu ,\mathrm{\Sigma }\right)$p(x^((i))|z^((i))=j;mu,Sigma)p\left(x^{(i)} | z^{(i)}=j ; \mu, \Sigma\right) is given by evaluating the density of a Gaussian with mean ${\mu }_j$${\mu }_j$mu_(j)\mu_{j} and covariance ${\mathrm{\Sigma }}_j$${\mathrm{\Sigma }}_j$Sigma_(j)\Sigma_{j} at $x^{(i)};p(z^{(i)}=j;\varphi )$$x^{(i)};p\left(z^{(i)}=j;\varphi \right)$x^((i));p(z^((i))=j;phi)x^{(i)} ; p\left(z^{(i)}=j ; \phi\right) is given by ${\varphi }_j$${\varphi }_j$phi_(j)\phi_{j}, and so on. The values $w_j^{(i)}$$w_j^{(i)}$w_(j)^((i))w_{j}^{(i)} calculated in the E-step represent our "soft" guesses^[2] for the values of $z^{(i)}$$z^{(i)}$z^((i))z^{(i)}.

Also, you should contrast the updates in the M-step with the formulas we had when the $z^{(i)}$$z^{(i)}$z^((i))z^{(i)}'s were known exactly. They are identical, except that instead of the indicator functions "$1\{z^{(i)}=j\}$$1\left\{z^{(i)}=j\right\}$1{z^((i))=j}1\left\{z^{(i)}=j\right\}" indicating from which Gaussian each datapoint had come, we now instead have the $w_j^{(i)}$$w_j^{(i)}$w_(j)^((i))w_{j}^{(i)}'s.

The EM-algorithm is also reminiscent of the K-means clustering algorithm, except that instead of the "hard" cluster assignments $c(i)$$c(i)$c(i)c(i), we instead have the "soft" assignments $w_j^{(i)}$$w_j^{(i)}$w_(j)^((i))w_{j}^{(i)}. Similar to K-means, it is also susceptible to local optima, so reinitializing at several different initial parameters may be a good idea.

It's clear that the EM algorithm has a very natural interpretation of repeatedly trying to guess the unknown $z^{(i)}$$z^{(i)}$z^((i))z^{(i)}'s; but how did it come about, and can we make any guarantees about it, such as regarding its convergence? In the next set of notes, we will describe a more general view of EM, one that will allow us to easily apply it to other estimation problems in which there are also latent variables, and which will allow us to give a convergence guarantee.