Deriving EM for Gaussian Mixture Models

1. The Gaussian Mixture Model and the Challenge of MLE

A Gaussian Mixture Model assumes that the data is generated from a weighted sum of $K$ different Gaussian distributions. The probability density function of a data point $\mathbf{x}$ is given by:

Here, the parameters $\theta =\{{\pi }_k,{\mathbf{\mu }}_k,{\mathbf{\Sigma }}_k{\}}_{k=1}^K$ are:

• ${\pi }_k$: The mixing coefficients, which are the prior probabilities of a data point coming from the $k$-th Gaussian. They must satisfy ${\pi }_k\ge 0$ and ${\sum }_{k=1}^K{\pi }_k=1$.
• ${\mathbf{\mu }}_k$: The mean vector of the $k$-th Gaussian.
• ${\mathbf{\Sigma }}_k$: The covariance matrix of the $k$-th Gaussian.

Our goal is to find the parameters $\theta$ that best fit the observed data $\mathbf{X}=\{{\mathbf{x}}_1,\dots ,{\mathbf{x}}_N\}$. The natural way to do this is through Maximum Likelihood Estimation (MLE). The log-likelihood function is:

If you were to try and maximize this by taking derivatives with respect to the parameters and setting them to zero, you would quickly run into a problem. The presence of the sum inside the logarithm prevents a closed-form solution. This is the central challenge that the EM algorithm is designed to solve.

2. The EM Algorithm: A Principled Framework for Latent Variables

The key insight of the EM algorithm is to reframe the problem by introducing latent variables. For each data point ${\mathbf{x}}_n$, we introduce a binary latent variable ${\mathbf{z}}_n$ of length $K$. If the data point ${\mathbf{x}}_n$ was generated by the $k$-th Gaussian, then the $k$-th element of ${\mathbf{z}}_n$, denoted $z_{nk}$, is 1, and all other elements are 0.

The problem is that we don't observe these $z_{nk}$ values. EM is a powerful iterative algorithm that allows us to find MLE solutions even in the presence of such hidden variables. It does so by repeatedly constructing a tight lower bound on the log-likelihood and then maximizing that bound. Your request for formal proofs makes the theoretical underpinnings of EM particularly relevant.

To understand why this iterative process is guaranteed to improve the likelihood at each step, we can turn to a rigorous derivation using Jensen's inequality.

As the video demonstrates, the EM algorithm consists of two repeating steps:

1. E-Step (Expectation): We don't know the true values of the latent variables $\mathbf{Z}$. Instead, we use our current parameter estimates, ${\theta }^{(t)}$, to compute the posterior distribution of these latent variables, $p(\mathbf{Z}|\mathbf{X},{\theta }^{(t)})$. We then use this to calculate the expectation of the complete-data log-likelihood. This expected value is called the Q-function, $Q(\theta |{\theta }^{(t)})$.

2. M-Step (Maximization): We find the new parameter estimate, ${\theta }^{(t+1)}$, by maximizing the Q-function with respect to $\theta$:

   $${\theta }^{(t+1)}=\arg \underset{\theta }{\max }Q(\theta |{\theta }^{(t)})$$

By iterating these two steps, we are guaranteed to find a local maximum of the original log-likelihood.

3. Deriving the E and M Steps for GMM

Now, let's apply this general framework to our GMM problem. This involves deriving the specific update equations for ${\pi }_k$, ${\mathbf{\mu }}_k$, and ${\mathbf{\Sigma }}_k$.

The E-Step: Calculating Responsibilities

In the context of GMM, the posterior probability of the latent variable, $p(z_{nk}=1|{\mathbf{x}}_n,{\theta }^{(t)})$, has a very intuitive name: the responsibility that component $k$ takes for data point ${\mathbf{x}}_n$. We'll denote it as ${\gamma }_{nk}$. Using Bayes' theorem, we can express it as:

Substituting the GMM definitions ($p(z_{nk}=1)={\pi }_k$ and $p({\mathbf{x}}_n|z_{nk}=1)=\mathcal{N}({\mathbf{x}}_n|{\mathbf{\mu }}_k,{\mathbf{\Sigma }}_k)$), we get the E-step update rule:

This is the "soft assignment" step. Instead of assigning ${\mathbf{x}}_n$ to a single cluster, we calculate its probability of belonging to each one. The following video segment walks through this derivation clearly.

The M-Step: Updating the Parameters

With the responsibilities ${\gamma }_{nk}$ calculated, we can now perform the M-step. We need to maximize the Q-function, which for GMM is:

We can maximize this with respect to ${\pi }_k$, ${\mathbf{\mu }}_k$, and ${\mathbf{\Sigma }}_k$ separately. The following video provides a detailed walkthrough of the necessary calculus. Since you're interested in the derivations, this is the core of the lesson.

Summary of M-Step Derivations

After performing the maximization as shown in the video, we arrive at the following update equations. Let's define the total responsibility for a cluster $k$ as $N_k={\sum }_{n=1}^N{\gamma }_{nk}$.

1. Update for the Means ${\mathbf{\mu }}_k$:
   Taking $\frac{\partial Q}{\partial {\mathbf{\mu }}_k}=0$ gives:

   $${\mathbf{\mu }}_k^{(t+1)}=\frac{1}{N_k}\sum _{n=1}^N{\gamma }_{nk}{\mathbf{x}}_n$$

   Intuition: The new mean for a cluster is a weighted average of all data points, where the weights are the responsibilities. Points that are highly likely to belong to the cluster contribute more to its new center.

2. Update for the Covariances ${\mathbf{\Sigma }}_k$:
   Taking $\frac{\partial Q}{\partial {\mathbf{\Sigma }}_k}=0$ gives:

   $${\mathbf{\Sigma }}_k^{(t+1)}=\frac{1}{N_k}\sum _{n=1}^N{\gamma }_{nk}({\mathbf{x}}_n-{\mathbf{\mu }}_k^{(t+1)})({\mathbf{x}}_n-{\mathbf{\mu }}_k^{(t+1)}{)}^T$$

   Intuition: The new covariance is a weighted covariance calculation. Each data point contributes to the covariance matrix of a cluster based on its responsibility for that cluster. Note that we use the newly computed mean ${\mathbf{\mu }}_k^{(t+1)}$ here.

3. Update for the Mixing Coefficients ${\pi }_k$:
   Maximizing with respect to ${\pi }_k$ subject to the constraint ${\sum }_k{\pi }_k=1$ (using a Lagrange multiplier) gives:

   $${\pi }_k^{(t+1)}=\frac{N_k}{N}$$

   Intuition: The new mixing coefficient for a cluster is simply the average responsibility that the cluster takes for all the data points.

For a consolidated written reference of these derivations, you can consult the following tutorial.

4. The Complete Algorithm

We can now summarize the entire EM algorithm for GMMs:

1. Initialize the means ${\mathbf{\mu }}_k$, covariances ${\mathbf{\Sigma }}_k$, and mixing coefficients ${\pi }_k$. This can be done randomly or by using the output of a simpler algorithm like K-means.

2. Iterate until convergence:
   a. E-Step: Evaluate the responsibilities ${\gamma }_{nk}$ using the current parameters ${\theta }^{(t)}$.

   $${\gamma }_{nk}=\frac{{\pi }_k^{(t)}\mathcal{N}({\mathbf{x}}_n|{\mathbf{\mu }}_k^{(t)},{\mathbf{\Sigma }}_k^{(t)})}{{\sum }_{j=1}^K{\pi }_j^{(t)}\mathcal{N}({\mathbf{x}}_n|{\mathbf{\mu }}_j^{(t)},{\mathbf{\Sigma }}_j^{(t)})}$$

   b. M-Step: Re-estimate the parameters ${\theta }^{(t+1)}$ using the responsibilities calculated in the E-step.
   \begin{align*}
   N_k &= \sum_{n=1}^N \gamma_{nk} \
   \boldsymbol{\mu}k^{(t+1)} &= \frac{1}{N_k} \sum{n=1}^N \gamma_{nk} \mathbf{x}n \
   \boldsymbol{\Sigma}k^{(t+1)} &= \frac{1}{N_k} \sum{n=1}^N \gamma{nk} (\mathbf{x}_n - \boldsymbol{\mu}_k^{(t+1)})(\mathbf{x}_n - \boldsymbol{\mu}_k^{(t+1)})^T \
   \pi_k^{(t+1)} &= \frac{N_k}{N}
   \end{align*}
   c. Check for convergence: Evaluate the log-likelihood and stop if the change between iterations is below a certain threshold.

Conclusion

In this lesson, we tackled the core mathematics behind Gaussian Mixture Models. We saw that direct maximization of the likelihood is intractable due to the sum inside the logarithm. This motivated the use of the Expectation-Maximization algorithm, a powerful general-purpose tool for MLE with latent variables.

Key Takeaways:

• The EM algorithm for GMMs is an iterative procedure that finds a local maximum of the data log-likelihood.
• The E-step is a "soft assignment" step where we compute the responsibilities ${\gamma }_{nk}$, which represent the posterior probability that data point ${\mathbf{x}}_n$ belongs to cluster $k$.
• The M-step is a maximization step where we update the model parameters (${\pi }_k,{\mathbf{\mu }}_k,{\mathbf{\Sigma }}_k$) by maximizing the expected complete-data log-likelihood.
• The update rules have clear intuitive interpretations: the new means and covariances are responsibility-weighted averages, and the new mixing coefficients are the average responsibilities per cluster.

Preview of the Next Lesson

We have now derived the update rules for both K-means (hard clustering) and GMMs (soft clustering). In the next lesson, we will step back from the derivations and focus on the practical implications. We will compare and contrast the use cases for K-means vs. GMM, discussing their respective strengths, weaknesses, and underlying assumptions, so you can make informed decisions about which clustering algorithm to apply to a given problem.