Expectation Maximization
EM is a broad topic. I hope to just give a brief but understandable introduction to it here. This material is borrowed heavily from a class I took, taught by Mathias Drton.
The setup
First, let's assume we have some complete data \(x\) , for which we know how to do Maximum Likelihood Estimation for some parameters \(\theta\). The complete data likelihood is then:
\[P(x \vert \theta)\]Now assume we don't observe this wonderful complete data, but only some transformation \(y = h(x)\). Then, the observed data likelihood is computed by:
\[P(y \vert \theta) = \sum_{x: y = h(x)} P(x \vert \theta)\]Let's take a very simple running example. Let's say we have some data \(x = [x_1,x_2,x_3, x_4, x_5]\) such that:
\[P(x | \pi) = \frac{(x_1 + x_2 + x_3 + x_4 + x_5)!}{x_1!x_2!x_3!x_4!x_5!}\left(\frac{1}{2}\right)^{x_1}\left(\frac{\pi}{4}\right)^{x_2}\left(\frac{1-\pi}{4}\right)^{x_3}\left(\frac{1-\pi}{4}\right)^{x_4} \left(\frac{\pi}{4}\right)^{x_5}\]That is - \(x\) is distributed according to a multinomial distribution with parameters \(\left[ \dfrac{1}{2}, \dfrac{\pi}{4}, \dfrac{1 - \pi}{4}, \dfrac{1- \pi}{4} , \dfrac{\pi}{4} \right]\)
We know that estimating \(\pi\) in this case is simple. We just take the derivative of the log likelihood and set it equal to 0:
\[\begin{align} log(P(x | \pi)) &\propto x_1 log(\frac{1}{2}) + (x_2 + x_5) log(\frac{\pi}{4}) + (x_3 + x_4) log(\frac{1 - \pi}{4}) \\ \frac{dlog(P(x | \pi))}{d\pi} &= \frac{(x_2 + x_5)}{\pi} - \frac{x_3 + x_4}{1 - \pi} = 0\\ \hat{\pi} &= \frac{(x_2 + x_5)}{x_2 + x_3 + x_4 + x_5} \end{align}\]Now, let's say that instead of observing \(x\), we observe \(y = [y_1,y_2,y_3,y_4]\), such that:
\[\begin{align*} y_1 &= x_1 + x_2 \\ y_2 &= x_3 \\ y_3 &= x_4 \\ y_4 &= x_5 \end{align*}\]Now, the likelihood \(P(y | \pi)\) is not so easy to optimize in close form (try it):
\[P(y | \pi) = \frac{(y_1 + y_2 + y_3 + y_4)!}{y_1!y_2!y_3!y_4}\left(\frac{1}{2} + \frac{\pi}{4}\right)^{y_1}\left(\frac{1-\pi}{4}\right)^{y_2}\left(\frac{1-\pi}{4}\right)^{y_3}\left(\frac{\pi}{4}\right)^{y_4}\]The algorithm
\(\require{color} \definecolor{red}{RGB}{255,0,0} \require{cancel}\) The EM algorithm works as follows:
- E-step - Evaluate \(E[log P(x \vert \theta) \vert y, \theta_n)]\)
- M-step - Set \(\theta_{n+1} = argmax_{\theta}E[log(P(x \vert \theta) \vert y, \theta_n)]\)
- repeat until some convergence criteria is satisfied
(e.g. \(\vert log P(y \vert \theta_{n+1}) - log P(y \vert \theta_n) \vert \le \epsilon\))
Now, what we want to prove is that the log-likelihood does not decrease at each iteration.
\[log P(y | \theta_{n+1}) \geq log P(y | \theta_n)\]Note that since \(y\) is a deterministic function of \(x\),
\[P(x | \theta_n) = P(x, y | \theta_n) = P(x | y, \theta_n)P(y | \theta_n)\]Or equivalently, we have:
\[P(y | \theta_n) = \frac{P(x | \theta_n)}{P(x | y, \theta_n)}\]Thus, we can rewrite the claim as:
\[log P(y | \theta_{n+1}) - log P(y | \theta_n) = log P(x | \theta_{n+1}) - log P(x | \theta_{n}) + log P(x | y, \theta_n) - log P(x | y, \theta_{n+1})\]Taking expectations \(E[. \vert y, \theta_n]\) over \(x\) on both sides:
\[log P(y | \theta_{n+1}) - log P(y | \theta_n) = \underbrace{ {\color{red} E[log P(x | \theta_{n+1})] - E[log P(x | \theta_{n})]}}_{\geq 0} + E[log P(x | y, \theta_n)] - E[log P(x | y, \theta_{n+1})]\]The red part is \(\geq\) 0 because of the way we choose \(\theta_{n+1}\). Note that due to the M-step, the left term has to be greater or equal to the right term. We can thus write the following inequality:
\[\begin{align}log P(y | \theta_{n+1}) - log P(y | \theta_n) &\geq E[log P(x | y, \theta_n)] - E[log P(x | y, \theta_{n+1})] \\ &= E\left[ -log\left(\frac{P(x | y, \theta_{n+1})}{P(x | y, \theta_{n})}\right) | y, \theta_n\right] \end{align}\]Now we use Jensen's inequality, which states that for a convex function \(g\):
\[E[g(x)] \geq g[E[x]]\]Since log is a concave function, -log is a convex function, and thus we can write:
\[\begin{align}log P(y | \theta_{n+1}) - log P(y | \theta_n) &\geq -log \left( E\left[\left(\frac{P(x | y, \theta_{n+1})}{P(x | y, \theta_{n})}\right) | y, \theta_n\right]\right) \\ &= -log \left( \sum_x \frac{P(x | y, \theta_{n+1})}{\cancel{P(x | y, \theta_{n})}}\cancel{P(x | y , \theta_n)}\right) \\ &= -log(1) \\ &= 0 \end{align}\]Thus, we have proved that the log likelihood does not decrease at each iteration.
Going back to the example
In the E step, we would evaluate the expected value of \(x_2, x_3, x_4, x_5\) given \(\pi_k\) (a current estimate of \(\pi\)) and \(\boldsymbol{y}\). Let's say each observation \(g_i\) in the complete data has a category in \(\{1,2,3,4,5\}\) and each observation \(f_i\) in the observed data has a category in \(\{1,2,3,4\}\) Since we assume data is i.i.d, \(P(g_i = 2 | \cdot) = P(g_1 = 2 | \cdot)\).
\[\begin{align*} E[x_2 | \pi_k, \boldsymbol{y}] &= E[\sum_{i = 1}^{n}{1_{\{g_i=2\}}} | \pi_k, \boldsymbol{y}] \\ &= \sum_{i=1}^{n}{P(g_i=2|\pi_k, \boldsymbol{y})} \\ &= y_1 * P(g_1=2|\pi_k, f_1 =1) \tag{*}\\ &= y_1 * \frac{P(g_1=2,f_1 = 1 | \pi_k)}{P(f_1=1 | \pi_k)} \\ &= y_1 * \frac{\pi_k/4}{1/2 + \pi_k/4}\\ &= \frac{y_1\pi_k}{2 + \pi_k}\\ E[x_3 | \pi_k, \boldsymbol{y}] &= y_2\\ E[x_4 | \pi_k, \boldsymbol{y}] &= y_3\\ E[x_5 | \pi_k, \boldsymbol{y}] &= y_4 \end{align*}\]For the \((*)\) line, note that \(P(g_1 = 2 \vert f_i \ne 1) = 0\) , and \(\sum_{i=1}^{n}\{f_i = 1\} = y_1\).
In the M step, we would maximize the expected log likelihood of the complete data using the values we found in the E step:
\[\begin{align*} E[log(P(\boldsymbol{x} | \pi)) | \pi_k, \boldsymbol{y}] &\propto (E[x_2|\pi_k, \boldsymbol{y}] + E[x_5|\pi_k, \boldsymbol{y}])log\left(\frac{\pi}{4}\right) + (E[x_3|\pi_k, \boldsymbol{y}] + E[x_4|\pi_k, \boldsymbol{y}]) log\left(\frac{1-\pi}{4}\right) \\ &\propto \left(\frac{y_1\pi_k}{2 + \pi_k} + y_4\right)log\left(\frac{\pi}{4}\right) + (y_2 + y_3) log\left(\frac{1-\pi}{4}\right) \end{align*}\]Taking the derivative and setting it to 0 gives us:
\[\begin{align*} \hat{\pi} = \frac{y_1\pi_k + y_4(2 + \pi_k)}{y_1\pi_k + (y_2 + y_3 + y_4)(2 + \pi_k)} \end{align*}\]Thus, we would set \(\pi_{k+1} = \hat{\pi}\) and iterate. Note that the equation for \(\hat{\pi}\) is exactly what we had for the complete likelihood, but taking expectations given the current parameter instead of using the values:
\[\hat{\pi} = \frac{(E[x_2 | \pi_k, y] + E[x_5 | \pi_k, y])}{E[x_2 | \pi_k, y] + E[x_3 | \pi_k, y] + E[x_4 | \pi_k, y] + E[x_5 | \pi_k, y]}\]This is an illustration of when EM is used: when the observed data likelihood is annoying, but the complete data likelihood is nice.