From Maximum Likelihood to EM Algorithm: A Consistent Understanding

Reasonable Existence

What is Maximum Likelihood? There is a philosophical saying: "Existence is reasonable." Maximum Likelihood means "Existence is the most reasonable." Specifically, if the probability distribution of an event X is p(X), and the values observed in a single observation are X_1, X_2, \dots, X_n, assuming they are independent and identically distributed (i.i.d.), then: \mathcal{P} = \prod_{i=1}^n p(X_i) \tag{1} should be maximized. If p(X) is a probability distribution p_{\theta}(X) with a parameter \theta, then we should find a way to choose \theta that maximizes \mathcal{P}(\theta), i.e., \theta = \mathop{\text{argmax}}_{\theta} \mathcal{P}(\theta) = \mathop{\text{argmax}}_{\theta}\prod_{i=1}^n p_{\theta}(X_i) \tag{2} Taking the logarithm of the probability yields the equivalent form: \theta = \mathop{\text{argmax}}_{\theta}\sum_{i=1}^n \log p_{\theta}(X_i) \tag{3} If we divide the right 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 refer to -\mathcal{L}(\theta) as the cross-entropy.

Theoretical Form

Theoretically, based on existing data, we can obtain the statistical frequency \tilde{p}(X) for each X. Thus, we can obtain an equivalent form of the above equation: \theta = \mathop{\text{argmax}}_{\theta} \mathcal{L}(\theta) = \mathop{\text{argmax}}_{\theta}\sum_X \tilde{p}(X)\log p_{\theta}(X) \tag{5} However, practically, it is almost impossible to obtain \tilde{p}(X) (especially for continuous distributions). What we can directly calculate is its mathematical expectation, which is equation (4), because calculating the expectation only requires computing the value for each sample, summing them up, and dividing by n. Therefore, equation (5) has primarily theoretical value, as it facilitates subsequent derivations.

It should be noted that the description above is very general. X can be any object; it could be continuous real numbers, in which case the summation is replaced by an integral and p(X) becomes a probability density function. Naturally, this presents no fundamental difficulty.

Broader KL Divergence

The form of maximum likelihood can also be derived from KL divergence. Suppose we have two distributions \tilde{p}(X) and p(X); 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, it is greater than 0. I assume the reader is already familiar with these properties.

Assuming the samples of X are given, this means \tilde{p}(X) can be regarded as known. At this point: \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 equations (4) and (5). In fact, KL divergence is richer in meaning than simple maximum likelihood because maximum likelihood assumes \tilde{p}(X) is known (given samples of X), which is not always achievable (e.g., in the context of the EM algorithm). Often, we only know partial information about X, at which point we must return to KL divergence.

Note: If the reader does not fully understand sampling calculation, please read the section "Numerical Calculation vs. Sampling Calculation" in "Variational Autoencoders (2): A Bayesian Perspective".

Classification Problems

Taking classification as an example, we usually model p(Y|X) rather than p(X, Y). That is, we want to determine the distribution of the output based on the input, rather than their joint distribution. Therefore, we still start from equation (5), using p(X, Y) = p(X)p(Y|X), to first get: \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 assume p_{\theta}(X) is simply \tilde{p}(X). This is equivalent to adding a constant term to the optimization objective, so equation (9) is equivalent to: \theta = \mathop{\text{argmax}}_{\theta} \sum_{X,Y} \tilde{p}(X, Y)\log p_{\theta}(Y|X) \tag{10} Then, since \tilde{p}(X, Y) = \tilde{p}(X)\tilde{p}(Y|X), equation (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, do not forget that in classification problems within supervised learning, generally for a specific input X in the training data, there is only one category. Thus, \tilde{p}(Y_t|X) = 1 and others are 0, where Y_t is the target label for X. Therefore: \theta = \mathop{\text{argmax}}_{\theta} \sum_X \tilde{p}(X) \log p_{\theta}(Y_t|X) \tag{12} This is 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}

Transformations

In fact, the content above consists merely of identity transformations and might not seem particularly valuable on its own, especially since the result (the cross-entropy loss for classification) is something we have long been familiar with. Therefore, this section merely demonstrates how to start from the most primitive form of the maximum likelihood function and eventually implement it in a specific problem, allowing the reader to become familiar with this step-by-step transformation process.

GMM Model

Wait, you want to do classification without label data? Of course, it’s possible. Isn’t the Gaussian Mixture Model (GMM) exactly such a model? In GMM, it is assumed that: p_{\theta}(X,Y) = p_{\theta}(Y) p_{\theta}(X|Y) \tag{14} Note that it is p_{\theta}(Y) p_{\theta}(X|Y) rather than p_{\theta}(X) p_{\theta}(Y|X). The difference is that it is difficult to directly estimate p(X) or guess the form of p(Y|X). However, p(Y) and p(X|Y) are relatively easy because we usually assume Y represents a category, so p(Y) is just a finite vector, and p(X|Y) represents the distribution of objects within each category. Since these objects belong to the same category, they should look similar. Thus, GMM assumes it is a normal distribution. In this case, the assumptions are grounded; otherwise, if all data were mixed together, who would know which distribution to assume?

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

pLSA Model

Of course, latent variables are not exclusive to unsupervised learning; supervised learning can have them too. For example, we can set: p(Y|X)=\sum_{Z}p_{\theta}(Y|Z)p_{\theta}(Z|X) \tag{15} Here, an additional variable Z is introduced. Even if (X, Y) label pairs are given, there is no data for Z. It is an imaginary variable we created, making it a latent variable. pLSA is such a problem. That is, the complete data pairs should be in the form (X, Y, Z), but we only know partial samples like (X_1, Y_1), \dots, (X_n, Y_n).

The Bayesian School

Some readers might think creatively: Can the parameter \theta also be viewed as a latent variable? Congratulations, if you have reached this level of insight, you have entered the realm of Bayesian thinking. The Bayesian school believes that everything is random and follows some probability distribution, and the parameter \theta is no exception. Unfortunately, Bayesian probability theory is quite profound, and we cannot put it to use here yet. (Actually, more importantly, the author doesn’t understand it yet either  )

Joint KL Divergence

Let’s first look at how general tutorials handle problems with latent variables. Since latent variables are unobservable, the maximum likelihood of the marginal distribution (the distribution of the observed variables) is generally used as the objective function: \theta = \mathop{\text{argmax}}_{\theta}\sum_X \tilde{p}(X)\log\sum_Z p_{\theta}(X|Z)p_{\theta}(Z) \tag{16} as the goal for maximization.

This approach is not impossible, but it requires introducing significant mathematical knowledge to derive the EM algorithm, and a rigorous proof requires a lengthy derivation. In fact, one 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 use the marginal distribution maximum likelihood approach, we cannot intuitively understand the origin of the Q function.

Taking GMM as an example, first, we calculate 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{constant} \end{aligned} \tag{17} Although this process is somewhat long, it involves no convoluted transformations and is relatively easy to accept.

The EM Master Arrives

Reviewing the source of equation (17), we hope to find a set of distribution parameters \theta such that KL\Big(\tilde{p}(X,Y)\Big\Vert p_{\theta}(X,Y)\Big) is as small as possible. p_{\theta}(X,Y) is already given in the form p_{\theta}(X|Y)p_{\theta}(Y), and only the parameter \theta is unknown. However, in equation (17), \tilde{p}(Y|X) is also unknown, including its form.

At this point, the master speaks: Let’s treat it as known for now. Then \tilde{p}(Y|X) can be regarded as a constant, and we can calculate 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}

Having calculated the new \theta^{(r)}, we then treat p_{\theta}(X|Y) as known to 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, equation (19) can be solved analytically. The answer 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 part inside the square brackets of equation (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{constant} \end{aligned} So minimizing equation (19) is equivalent to minimizing KL\Big(\tilde{p}(Y|X)\Big\Vert p_{\theta^{(r)}}(Y|X)\Big). According to the properties of KL divergence, the optimal solution is clearly when the two distributions are 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)} This yields equation (20).

Since we cannot minimize (17) in one step, we now train it alternately: fix one part, maximize the other, and then switch. The EM algorithm is an alternating training method for complex objective functions!

Combining equations (18) and (20) constitutes the entire solution algorithm. Now look at equation (18): it has an E (Expectation) and an M (\mathop{\text{argmax}}), so let’s call it the EM algorithm. The expression being expected, we’ll call the Q function. Thus, the EM master emerges, and the Q function emerges, just like that...

Of course, the "E" in the EM algorithm originally means treating \sum\limits_{Y}\tilde{p}^{(r-1)}(Y|X)\log p_{\theta}(Y) p_{\theta}(X|Y) as taking the expectation over the latent variable Y. Here we are being a bit casual, as long as the conclusion is correct 

Does it feel sudden? Like nothing was done, yet the EM algorithm was explained in just a few sentences? Including the derivation?

What is actually happening?

For pLSA or other models with latent variables, the EM algorithm can be derived similarly. Compared to the EM algorithm derivations I can find currently, I believe the process above is quite concise. Although there was a lot of preparation, it was all just basic knowledge.

How was this achieved? Reviewing the whole process, we didn’t actually do much; we purely used KL divergence as a measure of difference between joint distributions and then performed alternating minimization on the KL divergence. This derivation, starting from the joint distribution, turned out to be much more direct and faster than starting from the maximum likelihood of the marginal distribution, which was a pleasant surprise.