From Maximum Likelihood to the EM Algorithm: A Unified Way of Understanding

从最大似然到EM算法：一致的理解方式 苏剑林 · 2018-03-15 · Mathematics probability optimization clustering

I've been recently thinking about some topics related to unsupervised learning in NLP and probabilistic graphical models, so I went back and organized a few parameter estimation methods. In deep learning, parameter estimation is one of the most basic steps — this is exactly what we call the model training process. To train a model you need a loss function, and if a reader hasn't systematically studied probability theory, the most natural loss function they'd think of is probably the mean squared error, which corresponds to what we call Euclidean distance. Theoretically, though, the best-matched loss for a probabilistic model should be the "cross-entropy" function, which comes from the maximum likelihood function in probability theory.

Maximum Likelihood

Existence Is Reasonable

What is maximum likelihood? There's a philosophical saying: "what exists is reasonable." Maximum likelihood means "what exists is the most reasonable." Specifically, if the probability distribution of event $X$ is $p(X)$, and the values actually observed in one observation are $X_1,X_2,\dots,X_n$, and assuming they're mutually independent, then

$$\mathcal{P} = \prod_{i=1}^n p(X_i)\tag{1}$$

is the largest. If $p(X)$ is a probability distribution $p_{\theta}(X)$ with parameter $\theta$, then we should try to choose $\theta$ so as to maximize $\mathcal{L}$, i.e.

$$\theta = \mathop{\text{argmax}}_{\theta} \mathcal{P}(\theta) = \mathop{\text{argmax}}_{\theta}\prod_{i=1}^n p_{\theta}(X_i)\tag{2}$$more

Taking the log of the probability gives the equivalent form

$$\theta = \mathop{\text{argmax}}_{\theta}\sum_{i=1}^n \log p_{\theta}(X_i)\tag{3}$$

If we further divide the right-hand side by $n$, we get a more refined expression:

$$\theta = \mathop{\text{argmax}}_{\theta} \mathcal{L}(\theta) = \mathop{\text{argmax}}_{\theta} \mathbb{E}\big[\log p_{\theta}(X_i)\big]\tag{4}$$

where we call $-\mathcal{L}(\theta)$ the cross-entropy.

The Theoretical Form

In theory, based on the available data, we could obtain the statistical frequency $\tilde{p}(X)$ of each $X$, and then get the equivalent form of the above expression:

$$\theta = \mathop{\text{argmax}}_{\theta} \mathcal{L}(\theta) = \mathop{\text{argmax}}_{\theta}\sum_X \tilde{p}(X)\log p_{\theta}(X)\tag{5}$$

But in practice we can almost never get $\tilde{p}(X)$ (especially for continuous distributions); what we can compute directly is its mathematical expectation, i.e. expression $(4)$, since computing an expectation just requires computing the value for each sample, summing them up, and dividing by $n$. So expression $(5)$ only has theoretical value — it's convenient for later derivations.

Note that the above description is fully general: $X$ can be an arbitrary object, and it could even be a continuous real number, in which case the sum needs to be replaced by an integral, and $p(X)$ becomes a probability density function. This poses no essential difficulty.

The More General KL Divergence

Maximum likelihood can also be derived starting from KL divergence. Suppose we have two distributions $\tilde{p}(X)$ and $p(X)$, and we use KL divergence to measure the distance between them:

$$\begin{aligned}KL\Big(\tilde{p}(X)\Big\Vert p(X)\Big) =& \sum_X \tilde{p}(X) \ln \frac{\tilde{p}(X)}{p(X)}\\ =&\mathbb{E}\left[\ln \frac{\tilde{p}(X)}{p(X)}\right]\end{aligned}\tag{6}$$

When the two distributions are identical, the KL divergence is 0; when they differ, the KL divergence is greater than 0 — I'll assume the reader already knows these properties.

Now suppose the samples of $X$ have already been given, which means $\tilde{p}(X)$ can be regarded as known. Then:

$$\begin{aligned}\theta =& \mathop{\text{argmin}}_{\theta} KL\Big(\tilde{p}(X)\Big\Vert p_{\theta}(X)\Big)\\ =& \mathop{\text{argmax}}_{\theta}\sum_X \tilde{p}(X)\log p_{\theta}(X)\\ =& \mathop{\text{argmax}}_{\theta}\mathbb{E}\big[\log p_{\theta}(X_i)\big]\end{aligned}\tag{7}$$

This re-derives $(4)$ and $(5)$. In fact, KL divergence carries richer meaning than plain maximum likelihood, because maximum likelihood is equivalent to assuming $\tilde{p}(X)$ is known (i.e., we have samples of $X$), which isn't always achievable (as in the EM algorithm setting) — often we only know partial information about $X$, in which case we need to fall back on KL divergence.

Note: if the reader doesn't fully follow the sampling-based computation, please read the section Numerical Computation vs. Sampling Computation in Variational Autoencoders (II): From a Bayesian Perspective.

Supervised Models

Now let's see how the above is applied in supervised learning. Suppose the input is $X$ and the label is $Y$; then $(X,Y)$ constitutes an event, so from $(4)$ we have

$$\theta = \mathop{\text{argmax}}_{\theta} \mathbb{E}_{X,Y}\big[\log p_{\theta}(X,Y)\big]\tag{8}$$

Here we've explicitly noted that this is the mathematical expectation taken over the entirety of $X,Y$, though this expression isn't yet very practical.

Classification

Take classification as an example. What we typically model is $p(Y|X)$ rather than $p(X,Y)$ — that is, we want to determine the distribution of the output given the input, rather than their joint distribution. So we still need to start from expression $(5)$, and using $p(X,Y)=p(X) p(Y|X)$, first obtain

$$\theta = \mathop{\text{argmax}}_{\theta} \sum_{X,Y} \tilde{p}(X,Y)\log \big[p_{\theta}(X)p_{\theta}(Y|X)\big]\tag{9}$$

Since we only model $p(Y|X)$, we can treat $p_{\theta}(X)$ as being $\tilde{p}(X)$; this is equivalent to adding a constant term to the optimization objective, so $(9)$ is equivalent to

$$\theta = \mathop{\text{argmax}}_{\theta} \sum_{X,Y} \tilde{p}(X, Y)\log p_{\theta}(Y|X)\tag{10}$$

Then, we also have $\tilde{p}(X,Y)=\tilde{p}(X)\tilde{p}(Y|X)$, so expression $(8)$ can be further transformed into

$$\theta = \mathop{\text{argmax}}_{\theta} \sum_X \tilde{p}(X) \sum_Y \tilde{p}(Y|X)\log p_{\theta}(Y|X)\tag{11}$$

Finally, don't forget that we're dealing with a classification problem in supervised learning — generally, for a given input $X$ in the training data there's only one class, so $\tilde{p}(Y_t|X)=1$, with the rest being 0, and $Y_t$ is exactly the target label for $X$, so

$$\theta = \mathop{\text{argmax}}_{\theta} \sum_X \tilde{p}(X) \log p_{\theta}(Y_t|X)\tag{12}$$

This is precisely the most common maximum likelihood function for classification problems:

$$\theta = \mathop{\text{argmax}}_{\theta} \mathbb{E}_{X}\big[\log p_{\theta}(Y_t|X)\big]\tag{13}$$

Transform After Transform

In fact, everything above is just a series of identity transformations, and arguably has no particularly deep value in itself — its result (i.e., the cross-entropy loss for classification) is something we've already used to the point of rote familiarity. So this section merely demonstrates how, starting from the most primitive form of the maximum likelihood function, one can eventually arrive at a concrete problem, so readers get comfortable with this step-by-step transformation process.

Latent Variables

Now comes the moment to show its real value: we're going to use it to give a direct derivation of the EM algorithm (this blog also offers another way of understanding it — see Gradient Descent and the EM Algorithm: Same Root, Same Lineage). The EM algorithm is generally split into an M-step and an E-step; the M-step is relatively easy to understand, but the hard part is why the $Q$ function in the E-step is constructed the way it is. Many tutorials don't explain this $Q$ function at all, and some tutorials offer an explanation based on Jensen's inequality, but I don't think any of these approaches really bring out the essence of the EM algorithm.

Generally speaking, the EM algorithm is used to optimize probabilistic problems involving latent variables. What is a latent variable? It's simple — let's return to the classification example above. Classification models $p(Y|X)$, which of course is equivalent to $p(X,Y)$, and we said the objective is the maximum likelihood function, giving expression $(8)$:

$$\theta = \mathop{\text{argmax}}_{\theta} \mathbb{E}_{X,Y}\big[\log p_{\theta}(X,Y)\big]\tag{8}$$

If labeled pairs of $(X,Y)$ are given, that's just an ordinary supervised learning problem. But what if only $X$ is given, without $Y$? In this case, $Y$ is called a latent variable — it exists, but we cannot observe it, hence "latent."

The GMM Model

Wait — you want to do classification without labeled data? Certainly possible — isn't that exactly what the GMM (Gaussian mixture model) does? In GMM we assume

$$p_{\theta}(X,Y) = p_{\theta}(Y) p_{\theta}(X|Y)\tag{14}$$

Note this is $p_{\theta}(Y) p_{\theta}(X|Y)$, not $p_{\theta}(X) p_{\theta}(Y|X)$ — the difference being that it's hard to directly estimate $p(X)$, and it's also hard to directly guess the form of $p(Y|X)$. In contrast, $p(Y)$ and $p(X|Y)$ are relatively easy, because we typically take $Y$ to mean the class, so $p(Y)$ is just a finite vector, and $p(X|Y)$ describes the distribution of objects within each class — since objects within the same class should look fairly similar, GMM assumes it's a normal distribution, and this assumption is well justified. If instead you mixed all the data together, who would know which distribution to assume?

In this setting, our complete data should be $(X,Y)$, but we don't have such paired samples $(X_1,Y_1),\dots,(X_n,Y_n)$ (otherwise it would degenerate into supervised learning) — we only know samples $X_1,\dots,X_n$ of $X$, which corresponds exactly to the scenario we described in the KL divergence section.

The pLSA Model

Of course, latent variables aren't exclusive to unsupervised learning — supervised learning can have them too. For example, we could set

$$p(Y|X)=\sum_{Z}p_{\theta}(Y|Z)p_{\theta}(Z|X)\tag{15}$$

Here there's an extra variable $Z$: even given labeled pairs $(X,Y)$, $Z$ still has no data attached to it — it's a variable we imagine, i.e., a latent variable. pLSA is exactly this kind of problem. In other words, the complete data pairs here should take the form $(X,Y,Z)$, but we only ever know the partial samples $(X_1,Y_1),\dots,(X_n,Y_n)$.

The Bayesian School

A reader might have a "wild idea": could the parameter $\theta$ also be regarded as a latent variable? Congratulations — if you've had that realization, you've already stepped into the mindset of the Bayesian school. Bayesians believe that everything is random, that everything follows some probability distribution, and the parameter $\theta$ is no exception. Unfortunately, Bayesian probability theory is quite deep, and we can't put it to use here yet. (Actually, more importantly — I don't understand it myself yet!)

The EM Algorithm

All right, enough preamble — let's get down to formally discussing the EM algorithm.

Joint KL Divergence

Let's first look at how the standard tutorial approach handles solving problems involving latent variables. Since latent variables aren't observable, the usual approach switches to using the marginal distribution (i.e., the distribution over the observed variable) for maximum likelihood as the objective function, i.e., maximizing

$$\theta = \mathop{\text{argmax}}_{\theta}\sum_X \tilde{p}(X)\log\sum_Z p_{\theta}(X|Z)p_{\theta}(Z)\tag{16}$$

This approach isn't wrong per se, but it requires bringing in quite a bit of extra mathematics to arrive at the EM algorithm, and rigorous proofs need fairly lengthy derivations. In fact, we can start from KL divergence and greatly simplify the derivation of the EM algorithm by analyzing the KL divergence of the joint probability distribution. If we instead took the route of maximum likelihood on the marginal distribution, we wouldn't be able to intuitively understand where that $Q$ function comes from.

Taking GMM as an example, let's first compute the KL divergence between $\tilde{p}(X,Y)$ and $p_{\theta}(X,Y)$:

$$\begin{aligned} &KL\Big(\tilde{p}(X,Y)\Big\Vert p_{\theta}(X,Y)\Big)\\ =& \sum_{X,Y}\tilde{p}(X,Y)\log \frac{\tilde{p}(X,Y)}{p_{\theta}(X,Y)}\\ =& \sum_{X}\tilde{p}(X)\sum_{Y} \tilde{p}(Y|X)\log \frac{\tilde{p}(Y|X)\tilde{p}(X)}{p_{\theta}(X|Y)p_{\theta}(Y)}\\ =& \mathbb{E}_X\left[\sum_{Y} \tilde{p}(Y|X)\log \frac{\tilde{p}(Y|X)\tilde{p}(X)}{p_{\theta}(X|Y)p_{\theta}(Y)}\right]\\ =& \mathbb{E}_X\left[\sum_{Y} \tilde{p}(Y|X)\log \frac{\tilde{p}(Y|X)}{p_{\theta}(X|Y)p_{\theta}(Y)}+\sum_{Y} \tilde{p}(Y|X)\log \tilde{p}(X)\right]\\ =& \mathbb{E}_X\left[\sum_{Y} \tilde{p}(Y|X)\log \frac{\tilde{p}(Y|X)}{p_{\theta}(X|Y)p_{\theta}(Y)}\right]+\mathbb{E}\left[\log \tilde{p}(X)\right]\\ =& \mathbb{E}_X\left[\sum_{Y} \tilde{p}(Y|X)\log \frac{\tilde{p}(Y|X)}{p_{\theta}(X|Y)p_{\theta}(Y)}\right]+\text{const} \end{aligned}\tag{17}$$

This process, while fairly long, involves no roundabout tricks and should be easy to accept.

Enter the EM Master

Let's revisit the origin of expression $(17)$: we want to find a set of distribution parameters $\theta$ that make $KL\Big(\tilde{p}(X,Y)\Big\Vert p_{\theta}(X,Y)\Big)$ as small as possible. We've already given $p_{\theta}(X,Y)$ in the form $p_{\theta}(X|Y)p_{\theta}(Y)$, with only the parameter $\theta$ unknown. But in expression $(17)$, $\tilde{p}(Y|X)$ is also unknown — including its very form.

At this point, the master steps in and says: let's just treat it as known for now. Then $\tilde{p}(Y|X)$ can be treated as a constant, and we can compute the parameter $\theta$:

$$\begin{aligned}\theta^{(r)} =& \mathop{\text{argmin}}_{\theta} \mathbb{E}_X\left[\sum_{Y} \tilde{p}^{(r-1)}(Y|X)\log \frac{\tilde{p}^{(r-1)}(Y|X)}{p_{\theta}(X|Y)p_{\theta}(Y)}\right]\\ =& \mathop{\text{argmax}}_{\theta} \mathbb{E}_X\left[\sum_{Y}\tilde{p}^{(r-1)}(Y|X)\log p_{\theta}(Y) p_{\theta}(X|Y)\right]\end{aligned}\tag{18}$$

Then, having computed the new $\theta^{(r)}$, we treat $p_{\theta}(X|Y)$ as known and solve for $\tilde{p}(Y|X)$:

$$\tilde{p}^{(r)}(Y|X) = \mathop{\text{argmin}}_{\tilde{p}(Y|X)} \mathbb{E}_X\left[\sum_{Y} \tilde{p}(Y|X)\log \frac{\tilde{p}(Y|X)}{p_{\theta^{(r)}}(X|Y)p_{\theta^{(r)}}(Y)}\right]\tag{19}$$

In fact, expression $(19)$ has a closed-form analytical solution, which is:

$$\tilde{p}^{(r)}(Y|X)=\frac{p_{\theta^{(r)}}(Y)p_{\theta^{(r)}}(X|Y)}{\sum\limits_Y p_{\theta^{(r)}}(Y)p_{\theta^{(r)}}(X|Y)}\tag{20}$$

Supplementary derivation: the term inside the brackets in expression $(19)$ can be rewritten as
$$\begin{aligned}&\sum_{Y} \tilde{p}(Y|X)\log \frac{\tilde{p}(Y|X)}{p_{\theta^{(r)}}(X,Y)}\\ > =&\sum_{Y} \tilde{p}(Y|X)\log \frac{\tilde{p}(Y|X)}{p_{\theta^{(r)}}(Y|X)} - \sum_{Y} \tilde{p}(Y|X)\log p_{\theta^{(r)}}(X)\\ > =& KL\Big(\tilde{p}(Y|X)\Big\Vert p_{\theta^{(r)}}(Y|X)\Big) - \text{const} > \end{aligned}$$
so minimizing $(19)$ is equivalent to minimizing $KL\Big(\tilde{p}(Y|X)\Big\Vert p_{\theta^{(r)}}(Y|X)\Big)$. By the properties of KL divergence, clearly the optimal solution is for the two distributions to be exactly identical, i.e.
$$\tilde{p}(Y|X) = p_{\theta^{(r)}}(Y|X) = \frac{p_{\theta^{(r)}}(Y)p_{\theta^{(r)}}(X|Y)}{\sum\limits_Y p_{\theta^{(r)}}(Y)p_{\theta^{(r)}}(X|Y)}$$
which gives us expression $(20)$.

Since we can't solve for the minimum of $(17)$ in one shot, we now train it alternately: fix one part, maximize the other, and then swap. The EM algorithm is exactly this alternating training method for a complex objective function!

Combining expressions $(18)$ and $(20)$ gives us the complete solution algorithm. Now look at expression $(18)$: there's an E (expectation), and there's an M ($\mathop{\text{argmax}}$) — let's just call it the EM algorithm, and let's call the expression that gets an E taken of it the $Q$ function. And that's how the EM master was born, and the $Q$ function too — just like that, however we please...

Of course, the original intent of the "E" in EM is that $\sum\limits_{Y}\tilde{p}^{(r-1)}(Y|X)\log p_{\theta}(Y) p_{\theta}(X|Y)$ is viewed as taking an expectation over the latent variable $Y$ — here we're being a bit loose with the framing, but as long as the conclusion is correct, that's fine.

Does it feel a bit sudden? As if we barely did anything, and the EM algorithm was explained in just a couple of sentences — derivation included?

What Are We Actually Doing

For pLSA or other models with latent variables, the EM algorithm can be derived similarly. Comparing this with the derivations of the EM algorithm I've found elsewhere, I believe the process above is already remarkably concise. Although there was a lot of setup earlier, it was all just background material.

So how was this achieved? Looking back over the whole process, we really didn't do much — we simply used KL divergence as the measure of discrepancy for the joint distribution, then alternately minimized the KL divergence. The derivation obtained this way turns out to be significantly faster than starting from maximum likelihood on the marginal distribution — a pleasant surprise.

A Unified Understanding

This post reflects my own reflections on the principle of maximum likelihood. The overall approach starts from the principle and form of maximum likelihood, and derives from it a number of results in supervised/unsupervised learning, with the hope of stringing together various related topics under one unified idea. In the end, I find the results quite satisfying — especially the EM algorithm part: from now on, one only needs to remember that everything ultimately comes down to maximum likelihood (or KL divergence) of the (joint) distribution, and there's no more need to memorize the form of the Q-function in the EM algorithm by rote.

Of course, some of the views in this article are just "what I think," so there may be inaccuracies — readers should judge for themselves. What I can guarantee, though, is that the results agree with existing ones. I welcome further discussion from readers~