EM Algorithm

I'd like to talk about something about EM algorithm in my understanding.

This post is mainly based on Richard Xu's machine learning course.

Gaussian Mixture Model

Gaussian Mixture Model (GMM) (k-mixture) is defined as: $$\begin{array}{}\text{(1)}& p(X|\mathrm{\Theta })=\sum _{l=1}^k{\alpha }_l\mathcal{N}(X|{\mu }_l,{\mathrm{\Sigma }}_l)\end{array}$$

$$\begin{array}{}\text{(2)}& \sum _{l=1}^k{\alpha }_l=1\end{array}$$

and $$\begin{array}{}\text{(3)}& \mathrm{\Theta }=\{{\alpha }_1,\dots ,{\alpha }_k,{\mu }_1,\dots ,{\mu }_k,{\mathrm{\Sigma }}_1,\dots ,{\mathrm{\Sigma }}_k\}\end{array}$$

For data $X=\{x_1,\dots ,x_n\}$, we introduce latent variable $Z=\{z_1,\dots ,z_n\}$, each $z_i$ indicates which mixture components $x_i$ belongs to. (The introduction of latent variable should not change the marginal distribution of $p(X)$.)

Then we can use MLE to estimate $\mathrm{\Theta }$ : $$\begin{array}{}\text{(4)}& {\mathrm{\Theta }}_{MLE}=\underset{\mathrm{\Theta }}{\arg max}(\sum _{i=1}^N\log [\sum _{l=1}^k{\alpha }_l\mathcal{N}(x_i|{\mu }_l,{\mathrm{\Sigma }}_l)])\end{array}$$

This formula is difficult to solve because it is in 'log-of-sum' form. So, we solve this problem in an iterative way, called Expectation Maximization.

Expectation Maximization

Instead of performing $$\begin{array}{}\text{(5)}& {\mathrm{\Theta }}_{MLE}=\arg max\mathrm{\Theta }(\mathcal{L}(\mathrm{\Theta }))=\underset{\mathrm{\Theta }}{\arg max}(\log (p(X|\mathrm{\Theta })))\end{array}$$

we assume some latent variable $Z$ to the model, such that we generate a series of $\mathrm{\Theta }=\{{\mathrm{\Theta }}^{(1)},{\mathrm{\Theta }}^{(2)},\dots ,{\mathrm{\Theta }}^{(t)}\}$.

For each iteration of the E-M algorithm, we perform: $$\begin{array}{}\text{(6)}& {\mathrm{\Theta }}^{(g+1)}=\underset{\mathrm{\Theta }}{\arg max}({\int }_Z\log (p(X,Z|\mathrm{\Theta })p(Z|X,{\mathrm{\Theta }}^{(g)})))dZ\end{array}$$

We must ensure convergence: $$\begin{array}{}\text{(7)}& \log p(X|{\mathrm{\Theta }}^{(g+1)})\ge \log p(X|{\mathrm{\Theta }}^{(g)})\end{array}$$ Proof : $$\begin{array}{}\text{(8)}& E_{p(Z|X,{\mathrm{\Theta }}^{(g)})}[\log p(X|\mathrm{\Theta })]=E_{p(Z|X,{\mathrm{\Theta }}^{(g)})}[\log p(X,Z|\mathrm{\Theta })-\log p(Z|X,\mathrm{\Theta })]\end{array}$$

$$\begin{array}{}\text{(9)}& \log p(X|\mathrm{\Theta })={\int }_Z\log p(X,Z|\mathrm{\Theta })p(Z|X,{\mathrm{\Theta }}^{(g)})dZ-{\int }_Z\log p(Z|X,\mathrm{\Theta })p(Z|X,{\mathrm{\Theta }}^{(g)})dZ\end{array}$$

denote $$\begin{gathered} Q(\mathrm{\Theta },{\mathrm{\Theta }}^{(g)})={\int }_Z\log p(X,Z|\mathrm{\Theta })p(Z|X,{\mathrm{\Theta }}^{(g)})dZ \\ H(\mathrm{\Theta },{\mathrm{\Theta }}^{(g)})={\int }_Z\log p(Z|X,\mathrm{\Theta })p(Z|X,{\mathrm{\Theta }}^{(g)})dZ \end{gathered}$$ then we have $$\begin{array}{}\text{(10)}& \log p(X|\mathrm{\Theta })=Q(\mathrm{\Theta },{\mathrm{\Theta }}^{(g)})-H(\mathrm{\Theta },{\mathrm{\Theta }}^{(g)})\end{array}$$ Because $$\begin{gathered} Q({\mathrm{\Theta }}^{(g)},{\mathrm{\Theta }}^{(g)})\le Q({\mathrm{\Theta }}^{(g+1)},{\mathrm{\Theta }}^{(g)}) \\ H({\mathrm{\Theta }}^{(g)},{\mathrm{\Theta }}^{(g)})\ge H({\mathrm{\Theta }}^{(g+1)},{\mathrm{\Theta }}^{(g)}) \end{gathered}$$ the second inequality can be derived using Jensen's inequality.

Hence , $$\begin{array}{}\text{(11)}& \log p(X|{\mathrm{\Theta }}^{(g+1)})\ge \log p(X|{\mathrm{\Theta }}^{(g)})\end{array}$$

Using EM algorithm to solve GMM

Put GMM into this frame work. $$\begin{array}{}\text{(12)}& {\mathrm{\Theta }}^{(g+1)}=\underset{\mathrm{\Theta }}{\arg max}[Q(\mathrm{\Theta },{\mathrm{\Theta }}^{(g)})]=\underset{\mathrm{\Theta }}{\arg max}({\int }_Z\log (p(X,Z|\mathrm{\Theta })p(Z|X,{\mathrm{\Theta }}^{(g)})))dZ\end{array}$$ E-Step:

Define $p(X,Z|)$ : $$\begin{array}{}\text{(13)}& p(X,Z|\mathrm{\Theta })={\mathrm{\Pi }}_{i=1}^{n}p(x_i,z_i|\mathrm{\Theta })={\mathrm{\Pi }}_{i=1}^{n}p(x_i|z_i,\mathrm{\Theta })p(z_i|\mathrm{\Theta })={\mathrm{\Pi }}_{i=1}^n{\alpha }_{z_i}\mathcal{N}({\mu }_{z_i},{\mathrm{\Sigma }}_{z_i})\end{array}$$ Define $p(Z|X,\mathrm{\Theta })$ : $$\begin{array}{}\text{(14)}& p(Z|X,\mathrm{\Theta })={\mathrm{\Pi }}_{i=1}^{n}p(z_i|x_i,\mathrm{\Theta })={\mathrm{\Pi }}_{i=1}^n\frac{{\alpha }_{z_i}\mathcal{N}({\mu }_{z_i},{\mathrm{\Sigma }}_{z_i})}{\sum _{l=1}^k{\alpha }_l\mathcal{N}({\mu }_l,{\mathrm{\Sigma }}_l)}\end{array}$$ Then $$Q(\mathrm{\Theta },{\mathrm{\Theta }}^{(g)})=\sum _{z_1=1}^k\sum _{z_2=1}^k\dots \sum _{z_N=1}^k(\sum _{i=1}^N[\log {\alpha }_{z_i}+\log \mathcal{N}({\mu }_{z_i},{\mathrm{\Sigma }}_{z_i})]\ast {\mathrm{\Pi }}_{i=1}^{N}p(z_i|x_i,{\mathrm{\Theta }}^{(g)}))$$ $$\begin{array}{}\text{(15)}& =\sum _{i=1}^N\sum _{l=1}^k(\log {\alpha }_l+\log \mathcal{N}({\mu }_l,{\mathrm{\Sigma }}_l))p(l|x_i,{\mathrm{\Theta }}^{(g)})\end{array}$$

M-Step: $$\begin{array}{}\text{(16)}& Q(\mathrm{\Theta },{\mathrm{\Theta }}^{(g)})=\sum _{i=1}^N\sum _{l=1}^k\log ({\alpha }_l)p(l|x_i,{\mathrm{\Theta }}^{(g)})+\sum _{i=1}^N\sum _{l=1}^k\log \mathcal{N}({\mu }_l,{\mathrm{\Sigma }}_l)p(l|x_i,{\mathrm{\Theta }}^{(g)})\end{array}$$ The first term contains only $\alpha$ and the second term contains only $\mu ,\mathrm{\Sigma }$, so we can maximize both terms independantly.

Maximizing $\alpha$ means that: $$\begin{array}{}\text{(17)}& \frac{\partial \sum _{i=1}^N\sum _{l=1}^k\log ({\alpha }_l)p(l|x_i,{\mathrm{\Theta }}^{(g)})}{\partial {\alpha }_1\dots \partial {\alpha }_k}=0\end{array}$$ subject to $\sum _{l=1}^k=1$.

Solving this problem via Lagrangian Multiplier, we have $$\begin{array}{}\text{(18)}& {\alpha }_l=\frac{1}{N}\sum _{i=1}^{N}p(l|x_i,{\mathrm{\Theta }}^{(g)})\end{array}$$ Similarly, we can solve $\mu$ and $\mathrm{\Sigma }$.