Expectation-Maximization (EM) Algorithm: Solving a Chicken and Egg Problem
By: Matthew Prasad Burruss
Tags: Machine Learning, Data Science, Editors Pick, Statistics, Probability

In this task, the EM algorithm will be used to fit a Gaussian Mixture Model (GMM) to cluster the image into two segments. Our GMM will use a weighted sum of two (k=2) multivariate Gaussian distributions to describe each data point and assign it to the most likely distribution. The EM algorithm is used in this example to compute the parameters of the multivariate Gaussians distribution as well as the mixture weights.

Section 1. The Core Idea Behind Expectation-Maximization — Probabilistic Latent Representation

The main goal of expectation-maximization (EM) algorithm is to compute a latent representation of the data which captures useful, underlying features of the data. Using a probabilistic approach, the EM algorithm computes “soft” or probabilistic latent space representations of the data. For example, in Figure 1. each pixel is assigned a probability of being in class 0 and in class 1. As we will see later, these latent space representations in turn help us improve our understanding of the underlying statistical model, which in turn help us re-calculate the latent space representations, and so on.

The core goal of the EM algorithm is to alternate between improving the underlying statistical model and updating the latent representation of the data until a convergence criteria is met.

Commonly, the following notation is used when describing the EM algorithm and other related probabilistic models.

X: the raw data set where | X | = (# data points,# of features). X can be continuous or discrete (e.g. raw RGB values are continuous in the range [0,255]).

the raw data set where | | = (# data points,# of features). X can be continuous or discrete (e.g. raw RGB values are continuous in the range [0,255]). P(X|Z,θ): the complete likelihood of X under the underlying statistical model θ specified by Z where| P(X|Z,θ) | = (# data points, # of latent variables). The latent variable Z is discrete and specifies which statistical distribution X belongs to. In our example, X|Z,θ belongs to a multivariate Gaussian model and each pixel x can belong to either x|z=1,μ1,Σ1 or x|z=2,μ2,Σ2

the complete likelihood of under the underlying statistical model specified by where| | = (# data points, # of latent variables). The latent variable is discrete and specifies which statistical distribution belongs to. In our example, belongs to a multivariate Gaussian model and each pixel x can belong to either or P(Z=k): the prior probability of our latent representation Z or the probability of a randomly sampled data point taking on a particular latent assignment. Z typically belongs to a multinomial distribution or in our case the simpler binomial distribution.

the prior probability of our latent representation or the probability of a randomly sampled data point taking on a particular latent assignment. typically belongs to a multinomial distribution or in our case the simpler binomial distribution. θ: the parameters of our underlying statistical model. If we assume an underlying multivariate Gaussian mixture model, then θ includes both Σ and μ as well as the mixture coefficients π where |π| =(# of latent variables). The mixture coefficients represent the probability of randomly assigning a latent representation to a data point and is the same exact quantity as P(Z=k).

More clearly, in our example (assuming a GMM model) X is the 154401 x 3 pixel matrix, Z is the 154401 x 1 clustering assignments, the statistical model parameters θ = {μ1, Σ1, μ2, Σ2, π1, π2} where |μ_i| = 3, |Σ_i| = 3 x 3, and π1 + π2 = 1.

Section 2. Understanding Expectation-Maximization

Now that we have a concrete example to work with, let’s piece apart the definition of the EM algorithm as an “iterative method that updates unobserved latent space variables to find a local maximum likelihood estimate of the parameters of a statistical model” [2].

Part 1. Setup: “An iterative method…”

The EM algorithm has three main steps: the initialization step, the expectation step (E-step), and the maximization step (M-step). In the first step, the statistical model parameters θ are initialized randomly or by using a k-means approach. After initialization, the EM algorithm iterates between the E and M steps until convergence.

Part 2. “…that updates unobserved latent space variables…”

The E-step is used to update the unobserved latent space variables Z and set the stage for updating the parameters θ of the statistical model. The E-step can be broken down into two parts. The first part updates our conditional distribution P(Z|X,θ) which represents the soft latent assignments of our data (see Appendix “Soft Latent Assignments” for calculations). The second part of the E-step calculates Q(θ,θ*) where θ* represents the previous parameters of the statistical model and θ represents potential new parameters.

The current values of our statistical model θ* and the data X are used to compute soft latent assignments P(Z|X,θ*).

The intuition behind Q(θ,θ*) is probably the most confusing part of the EM algorithm. Its inclusion ultimately results in a trade-off between computation time and optimality. To explain, the disadvantage of the EM algorithm is that it is only guaranteed to find an estimate of θ that finds a local maximum of the likelihood P(X|θ) and not necessarily the absolute maximum. In a perfect world, the true maximum likelihood estimate could be found by maximizing the likelihood P(X,Z|θ); however, maximizing this value is difficult because we need to marginalize over Z in order to obtain the likelihood P(X|θ) (for more details, see Appendix “Intractability of Log Likelihood”). Instead, the EM algorithm maximizes Q(θ,θ*) which is related to P(X|θ) but is easier to optimize.

More specifically Q(θ,θ*) is the expectation of the complete log-likelihood log[P(X|Z,θ)] with respect to the current distribution of Z given X and the current estimates of θ*. In other words, it is the expectation of the complete log-likelihood with respect to the previously computed soft assignments Z|X,θ*.

To understand why we need Q(θ,θ*), think about this. We need to find the best θ to maximize P(X,Z|θ); however, we can’t reasonably sum across all of Z for each data point. Suppose however that Z was magically known. Then this problem could be avoided altogether because P(X,Z|θ) would become P(X|Z,θ). However, the obvious problem is Z is not known at the start. To get around this, we compute P(Z|X,θ*) to provide a soft estimate of Z and take the expectation of the complete log-likelihood conditioned on Z|X,θ* to fill in Z. In other words, we condition the expectation of P(X|Z,θ) on Z|X,θ* to provide a “best guess” at the parameters θ that maximize the likelihood P(X|Z,θ).

The simplified version of Q(θ,θ*) is shown below (see Appendix “Calculating Q(θ,θ*)” for details).

The E-step is used to find Q(θ,θ*), which is the expectation of the complete log-likelihood with respect to Z conditioned on the previous statistical model parameters θ* and the data X.

Part 3: “…to find a local maximum likelihood estimate (MLE) of the parameters of a statistical model.”

Compared to the E-step, the M-step is incredibly simple and is used to update the parameters θ of our statistical model. Once, we have computed Q(θ,θ*) we can improve the MLE of the statistical model parameters by evaluating the following expression:

The M-step operation to update the statistical model.

Part 4. Convergence!

The convergence criteria is simple — each new computation of Q(θ,θ*) is compared to the previous, and if the difference is less than some threshold ϵ (e.g. ϵ = 1e-4) the EM algorithm terminates. In practice, you would want to run the algorithm several times with various initializations of θ to find the parameters that most maximize P(X|Z,θ) because you are only guaranteed to find a local maximum likelihood estimate each time the EM algorithm executes.

Section 3. A Python Example Using EM to Parameterize a Gaussian Mixture Model

Code Snippet: An example of a Gaussian Mixture Model (GMM) which uses expectation maximization to cluster an image.

Figure 2. Example segmentations of images from the Berkeley Segmentation Data set [3] using EM to parameterize a Gaussian Mixture Model (k=2). Source: Photo by Author.

Section 4. Three Important Intuitions of EM

In summary, there are three important intuitions behind Q(θ,θ*) and ultimately the EM algorithm. First, the complete log-likelihood P(X|Z,θ) is faster to maximize than the log likelihood P(X,Z|θ) because there is no marginalization over Z. However, this introduces a problem because we don’t know Z. Therefore, the second intuition is that we can instead maximize Q(θ,θ*) or the expected value of the log of P(X,|Z,θ) where Z is filled in by conditioning the expectation on Z|X,θ*. The final intuition is that by finding the parameters θ that maximize Q(θ,θ*) we will be closer to a solution which would maximize the likelihood P(X,Z|θ). While non-trivial, the proof of this correctness shows that improving Q(θ,θ*) causes P(X,Z|θ) to improve by at least as much if not more. Thus, the EM algorithm will always converge to a local maximum.

Section 5. Recent Work and Conclusions

Expectation-Maximization (EM) algorithm originally described by Dempster, Laird, and Rubin [1] provides a guaranteed method to compute a local maximum likelihood estimation (MLE) of a statistical model that depends on unknown or unobserved data. Although it can be slow to execute when the data set is large; the guarantee of convergence and the algorithm’s ability to work in an unsupervised manner make it useful in a variety of tasks. More recently, there has been work on using neural networks, mainly an encoder-decoder architecture, where the encoder is used to re-parameterize X into high-level variables X* and the decoder encapsulates the statistical parameters θ* that produce the soft assignments P(Z|X*,θ*) per pixel. The network is trained using a loss function typical of encoder-decoders, but is weighted by P(Z|X*,θ*). This method is known as neural expectation maximization (N-EM) [4] and although useful, it loses the convergence guarantee of EM. Furthermore, it is unclear whether or not this approach is extracting more than just similarly colored features from images, leaving ample room for improvement and further study.

An exciting challenge in the field of AI will be developing methods to reliably extract discrete entities from raw sensory data which is at the core of human perception and combinatorial generalization [5]. Expectation-maximization, although nothing new, provides a lens through which future techniques seeking to develop solutions for this problem should look through.

Appendix

A. Soft Latent Assignment Calculations

The soft assignments are computed during the expectation step (E-step) to update our latent space representation.

Supplement 1: Soft assignments of latent space variable calculations derived using Bayes Rule and the Law of Total Probability.

B. Explaining the Intractability of Log-Likelihood

Typically, the optimal parameters of a statistical model are fit to data by finding θ which maximizes the log-likelihood or log[P(X|θ)]. Assuming independence, this typically looks like the following:

Equation 1: Typical maximum likelihood estimatation (MLE) setup.

However, we can only compute P(X,Z|θ) due to the dependency on Z and thus to compute P(X|θ) we must marginalize out Z and maximize the following:

Equation 2. The log-likelihood.

This quantity is more difficult to maximize because we have to marginalize or sum over the latent variable Z for all n data points.

Computing Q(θ,θ*)

Instead of maximizing the log-likelihood in Equation 2, the complete data log-likelihood is maximized below which at first assumes that for each data point x_i we have a known discrete latent assignment z_i.

Equation 3. The complete log-likelihood

The good news, unlike Equation 2. we no longer have to sum across Z in Equation 3. However, the bad news is that we don’t know z_i. To solve this, we try to guess at z_i by maximizing Q(θ,θ*) or the expectation of the complete log-likelihood with respect to Z|X,θ* which allows us to fill in the values of z_i.

Equation 4. The expected value of the complete log-likelihood with respect to the current conditional distribution Z|X,θ*.

Equation 4 can be simplified into the following

where I is the indicator function and can be used to evaluate the expectation because we assume that z_i is discrete. Ultimately, the equation above can simplify to

Equation 5. The simplified version of Q(θ,θ*)

Equation 5. finally shows the usefulness of Q(θ,θ*) because, unlike Equation 3, no terms in the summation are conditioned on both Z and θ. In the equation above, the left-most term is the soft latent assignments and the right-most term is the log product of the prior of Z and the conditional P.M.F. The right-most term can be separated into two terms allowing for the maximization of the mixture weights (prior of Z) and the distribution parameters of the P.M.F.

Going back to the concrete GMM example, while it may not be obvious above in Equation 5., {μ1, Σ1}, {μ2, Σ2},and {π1, π2} appear in different terms and can be maximized independently using the known MLEs of the respective distributions. For example, when updating {μ1, Σ1} and {μ2, Σ2} the MLEs for the Gaussian can be used and for {π1, π2} the MLEs for the binomial distribution. The only difference between these updates and the classic MLE equation is the inclusion of the weighting term P(Z|X,θ*).

Citations

[1] Dempster, A.P.; Laird, N.M.; Rubin, D.B. (1977). “Maximum Likelihood from Incomplete Data via the EM Algorithm”. Journal of the Royal Statistical Society, Series B. 39: 1–38.

[2] “Expectation-Maximization algorithm”, Wikipedia article, https://en.wikipedia.org/wiki/Expectation%E2%80%93maximization_algorithm

[3] Hui Li, Jianfei Cai, Thi Nhat Anh Nguyen, Jianmin Zheng. A BENCHMARK FOR SEMANTIC IMAGE SEGMENTATION. IEEE ICME 2013.

[4] Greff, Klaus, Sjoerd Van Steenkiste, and Jürgen Schmidhuber. “Neural expectation maximization.” Advances in Neural Information Processing Systems. 2017.

[5] Battaglia, Peter W., et al. “Relational inductive biases, deep learning, and graph networks.” arXiv preprint arXiv:1806.01261 (2018).