Skip to content

Expectation-Maximization (EM Algorithm)

Definition

The EM algorithm is based on iteratively applying two steps: E-Step and M-Step.

E-Step updates the probability array \(p_{t}[i, k]\), which states that the probability of sample \(i\) came from component \(k\), at \(t\) time step.

M-Step updates the model parameters by solving a minimization problem

The beauty of EM is that the component distributions are optimized independently of each other.


E-Step

To calculate E-Step, we use this formula:

\[ p_{t}[i,k] = Pr(k|\mathbf{x}_{i}) = \frac{P(\mathbf{x}_{i}|k)P(k)}{P(\mathbf{x}_{i})} = \frac{w_{k}^{(t)} \cdot p(\mathbf{x}_{i} | \boldsymbol{\theta}_{k}^{(t)})}{\sum_{j=1}^m w_{j}^{(t)}p(\mathbf{x}_{i}| \boldsymbol{\theta}_{j}^{(t)})} \]

Where:

\[ \begin{align} P(\mathbf{x}_{i}|k) &= w_{k}^{(t)} \\ P(k) &= p(\mathbf{x}_{i}|\boldsymbol{\theta}_{k}^{(t)}) \\ P(\mathbf{x}_{i}) &= \sum_{j=1}^m w_{j}^{(t)}p(\mathbf{x}_{i}|\boldsymbol{\theta}_{j}^{(t)}) \end{align} \]

M-Step

The minimization problems that we want to solve is:

\[ \begin{align} &\arg\!\min_{\mathbf{w} \in \mathcal{F}} -\sum_{k=1}^m p_{t}[k]\ln w_{k} \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ (1) \\ &\arg\!\min_{\boldsymbol{\theta}_{k}} - \sum_{i=1}^n p_{t}[i, k] \ln p(\mathbf{x}_{i} | \boldsymbol{\theta}_{k}) \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ (2) \end{align} \]

Equation \((1)\) measures the cross entropy between \(p_{t}[k]\), the posterior probability of seeing the \(k\)'th component across all samples, and \(w_{k}\), the weights that we learn corresponding to each component.

We want to find the vector \(\mathbf{w}\), the collection of all weights, that follows the simplex constraint \(\mathcal{F}\), such that we minimize the cross entropy loss between the two distributions

Equation \((2)\) measures the cross entropy between the distributions \(p_{t}[i,k]\), the posterior probability of the \(i\)'th sample coming from the \(k\)'th component, and \(\ln p(\mathbf{x}_{i}| \boldsymbol{\theta}_{k})\), the likelihood of getting \(\mathbf{x}_{i}\) by sampling from the \(k\)'th distribution with parameters \(\boldsymbol{\theta}_{k}\). In each iteration \(t\),

To solve for \((1)\) we convert it to an unconstrained optimization problem using the Lagrangian:

\[ L(w, a, b) = -\sum_{j=1}^m p_{t}[i, k] \ln(w_{k}) + a\left( \sum_{k=1}^m w_{k} - 1 \right) - b \sum_{j=1}^m b_{k}w_{k} \]

To solve for \((2)\), it is simple to solve log likelihood problems with simple distributions given by \(\boldsymbol{\theta}_{k}\).


Convergence Criteria

A common convergence criteria is to check if \(\sum_{j=1}^m p_{t}[k]\ln w_{k}\) has barely changed between iterations \(t\).

This criteria checks if each component \(k\) has converged, which is sufficient enough to say the whole optimization converged.

References

Cross-Entropy (Probability) Simple Generative Models, Mixture Models