In many settings, we have data that is generated by some underlying causes that we are actually interested in. Because we do not observe these causes, we cannot model them directly. Instead, we model them as latent variables and use the data we have to learn them or infer their parameters.

Fig. 1: Graphical representation of latent variable models

Consider a latent variable model: $Z \rightarrow X$ . Here, $X$ is our observed data, and we assume that it comes from the latent variable $Z$. In particular, the likelihood for $X$ is

that is we marginalize Z out. Denote $p(z) := p(Z=z| \theta)$ and $p( x|z, \theta) = p(X=x | Z = z, \theta)p(x|z, \theta) = p(X=x | Z = z, \theta)$ where $x \in \mathcal X$. Consider now a dataset $\mathcal D = \{ x_i \}_{i=1}^n$. Its log-likelihood is

We can see that the equation above is non-linear w.r.t. to the parameters of the model $\theta$; hence, we do not have a closed-form solution.

Expectation Maximization (EM)

To find the model's parameters $\theta$ , we instead use an iterative procedure. It consists of two steps, namely expectation and maximization.

First, we consider the complete log likelihood $\log p(\mathcal D, z | \theta)$. We then define the loss function

where $\theta_t$ is the parameters of the previous iteration. We know that

Combining $\mathcal L (\theta, \theta_t)$ and $p(\mathcal D, z| \theta)$ yields

To maximize $\mathcal L(\theta, \theta^t)$, we alternate between

1. Expectation step computing the posterior $p(z| \mathcal D, \theta^t), \forall z \in \mathcal Z$.
2. Maximization step finding the model's parameters $\theta$ that maximize $\mathcal L (\theta, \theta^t)$; this can be done by $\nabla_\theta \mathcal L(\theta, \theta^t) \overset{!}{=} 0$ and solve for each $\theta$.

It can be shown that the loss function increases every step; that is $\mathcal{L}(\theta, \theta^{t+1}) \ge \mathcal{L}(\theta, \theta^t)$.

Mixture Models

Mixture of Gaussians

We assume that our data $\mathcal D = \{ \mathbf x_i \in \reals^d \}_{i=1}^n$ is generated from some latent factor $z \in [1, K]$. In particular, each sample's likelihood is

where $\pi_k$'s are mixing weights and $\sum_k \pi_k = 1$. Here we assume that each component is a Gaussian distribution with $\theta_k= \{ \mu_k, \Sigma_k \}$ ; hence, we have

Consider $\mathbf z_i \in \{0, 1\}^K$ be the true mixture assignment of $\mathbf x_i$ (i.e., it is a one-hot vector whose is one at the cluster that $\mathbf x_i$ belongs to). The complete log likelihood is

Let $z_i$ be the cluster index for $\mathbf x_i$ (the index of the one entry). We compute the expected complete log likelihood w.r.t to the approximate posterior $q(z|\mathbf x, \theta^{t})$:

Define the responsibility (or posterior) $r_{ik} := p(z_i=k|\mathbf x_i, \theta^{t})$. We arrive at

Expectation Step

Due to the fact that

we can then write

Maximization Step

We compute $\forall _k, \pi_k, \theta_k := (\mu_k, \Sigma_k)$ that maximize $\mathcal{L}(\theta, \theta^{t})$.

Updating $\pi_k$

We know that $\sum_k \pi_k = 1$. Using Lagrange's multiplier, we get

Computing $\frac{\partial L}{\partial \pi_k}$ and set it to zero yields

Combing this into the constraint, we have

Hence, $\lambda = N$ and $\pi_k = \frac{1}{N} \sum_i r_{ik}$.

Updating $\mu_k$

Recall that $\mathcal{N}(\mathbf x| \mu_k, \Sigma_k) = \frac{1}{(\sqrt{(2\pi)^d|\Sigma_k|}) }\exp\bigg(-\frac{1}{2}(\mathbf x - \mu_k)^T\Sigma_k^{-1}(\mathbf x - \mu_k)\bigg)$; hence, we have

Taking the derivative w.r.t $\mu_k$ we get

Setting it to zero yields

Updating $\Sigma_k$

Taking the derivative yields

We also have that

Combining the two terms and set the derivative to zero, we arrive at

Therefore, we have

Fig. 2: Expectation-Maximization steps for a mixture of Gaussians

K-Means Algorithm: Special Case of GMMs

One can view K-means as a special case of GMMs that greedily chooses the closest centroid for each $\mathbf x_i$ . In this case, we assume that $\pi_k = \frac{1}{K}$ and $\Sigma_k = \sigma^2I_d$ for $\sigma \in \R_+$. The posterior is

where $z^*_i = \argmax_k p(z_i =k| \mathbf x_i, \theta)$. Equivalently, this is similar to

As you can see from the derivation, this is a hard assignment assigning a cluster to the sample. Consider $\mathcal D_k \subset \mathcal D$ be the samples belong to the $k$-th cluster. the maximization step is then just

Mixture of Binomials

Consider $z \in\{ H, T\}$ and we assume that it comes from a Bernoulli distribution:

where $\pi_H + \pi_T = 1$. Our setting has two Binomial mixture components, namely

where $x \in [0, m]$. In our setting, we have $\theta := \{ \pi_H, \pi_T, \mu_H, \mu_T\}$. Let $\mathcal{D} = (x_1, \dots, x_n).$ Recall the complete log-likelihood:

Expectation Step

We first recall and write

Similarly, we have

Maximization Step

Recall that $\theta = \{ \pi_H, \pi_T, \mu_H, \mu_T\}$. We compute the derivative of $\mathcal{L}( \theta, \theta^{t})$ w.r.t the parameters and set to zero.

Updating $\pi_H$ and $\pi_T$

We get

and $\pi_T = 1 - \pi_H$.

Updating $\mu_H$ and $\mu_T$

Therefore, we have

Similarly, we have $\mu_T = \frac{\sum_{i=1}^n {r_{iT}x_i} }{\sum_{i=1}^n {r_{iT}m} }$.

Fig. 3: Expectation-Maximization steps for a mixture of binomials

Conclusions

These are resources that I consulted while writing this article

• Murphy (2012), "Machine Learning: A Probabilistic Perspective" (Section 11.4) for the general idea of the EM algorithm.
• Manfred Opper's Probabilistic Bayesian Modeling course, Lecture Laplace Approximation (TU Berlin, Summer 2020) for Proof of the EM step
• Woijeck Samek's Machine Learning 1 course, Lecture Latent Variable Models (TU Berlin, Winter 2020) for the Mixture of Binomials example

This articles' figures are made in Google Colab.

Proof: the EM procedure increases the expected complete log-likelihood

In the following, we will show that $\mathcal{L}(\theta, \theta^{t+1}) \ge \mathcal{L}(\theta, \theta^t)$.

Consider $q(\cdot)$ be any distribution over the latent variable $Z$. We know that $p(\mathcal D, z| \theta) = p(z|\mathcal D, \theta) p(\mathcal D | \theta)$ . We can then compute

Because $\text{KL}(\cdot \| \cdot) \ge 0$, we have

Now consider $q_t(z) := p(z | \mathcal D, \theta^t)$. We have

Consider the fact that $F(q_t, \theta) \ge - \log p(\mathcal D| \theta)$. We have two cases, namely

1. $F(q_t, \theta) > - \log p(\mathcal D|\theta)$
2. $F(q_t, \theta) = - \log p( \mathcal D | \theta) \implies F(q_t, \theta_t) = - \log p(\mathcal D |\theta_t)$.

Subtract these two conditions lead to

In other words, we have

$(*)$ is due to the fact that we maximize $\mathcal{L}(\theta, \theta^t)$ w.r.t $\theta$, which is equivalent to minimize $F(q_t, \theta)$.

This concludes that the EM procedure always increases the expected complete log likelihood. However, the procedure does not guarantee to reach the global optimum.