Latent Variables & Expectation Maximization Algorithm
By: Saptashwa Bhattacharyya
Tags: Data Science, Towards Data Science, Bayesian Machine Learning, Clustering, Machine Learning

Variational Lower Bound and EM Algorithm:

It was just mentioned, until convergence repeat the steps of EM algorithm. You might think how to guarantee that the algorithm will converge? Let’s dig deeper! We started the GMM saying that for every data-point x_n there’s a latent variable z_n. The log likelihood of the complete data-set takes the form as below —

Exp 9: The complete likelihood including latent variables

As mentioned before, the problem with this expression is Log is acting on the sum of the mixtures of the Gaussian. We will play little trick to do away with this issue. Let’s introduce arbitrary probability distribution (without any special property) over latent variable — q(z) and now we can modify the expression above —

Exp 10: Little modification in the data likelihood

Using the definition of Expectation value this expression can be further simplified as —

Exp 11: Use the definition of expectation

The expectation term can be read as expected value of z_i drawn from a distribution q_i. We can compare this expression with Exp. 8 above to understand how summing over the latent variables gives us the expectation. In Exp. 11, Log is operating on the expectation value. Given that Log is a concave function, let’s apply Jensen’s inequality to further simplify the expression —

Exp 12: Use Jensen’s inequality to modify Exp 11.

The expression above says that log likelihood of the complete data-set is greater than or equal to the above expectation. The expectation term is known as Evidence Lower Bound (ELBO) and, expanding the logarithm and applying Bayes’ theorem we can reach more meaningful conclusion —

Exp 13: Evidence Lower Bound

Using the definition of KL divergence the above expression becomes —

Exp 14: ELBO and KL Divergence

As KL divergence ≥ 0, this means ELBO ≤ log p(x). Now we know why it is called a lower bound. Our objective is to maximize our data likelihood log p(X), and instead we maximize the ELBO. Based on this expression (Exp. 14) let’s review the steps of EM algorithm one more time —

Start the algorithm by randomly initializing the parameters. E step: Maximize the ELBO w.r.t q(z) and keep the parameters (π, μ, Σ) fixed. Since log p(X) is independent of q(z), this step is essentially just calculating the KL divergence and in case of GMM the q(z) is set equal to the posterior. This is possible since the calculation of posterior for every data-point is possible in GMM. M-step: ELBO is maximized w.r.t to the parameters while q(z) is kept fixed and this is very similar like MLE estimation. In this step the parameters are updated.

Why Convergence:

We started this section with a question- why EM algorithm ensures convergence? EM algorithm is an iterative process and thus E and M step goes on in cycle. It is important to remember that in each steps of EM algorithm, first the distribution q(z) is set equal to the posterior p(z|x) (E-step), and the parameters are updated in the M step from maximization. With the updated parameters the lower bound increases (unless the algorithm already reached a local maxima). This, in turn, increases the log likelihood but, logp(x) increases more than the lower bound because, while q(z) was obtained for old parameters but, now in E step when we calculate the posterior p(z|x) again and this time using the new parameters from M-step, gives rise to a non-zero KL divergence term. Thus the lower bound, as well as the data likelihood, increases monotonically. To ensure that the algorithm does not get stuck in a local maxima, we run the algorithm for different initial values of the parameters. This monotonic increasing nature of the data-likelihood is what we wanted to show as a proxy for convergence. This sometimes is very helpful to debug the code if we see the likelihood is not increasing in each step. To run the algorithm we will give the number of iteration as a parameter and also set a tolerance for the fractional increase in loss function. Now, it is time to put all the theory into action and let’s code!! Please check the notebook for detailed code along with other references in the reference section.