CS229 EM Algorithm & Gaussian Mixture Model Notes — Complete Derivation

These are complete CS229 EM algorithm and Gaussian Mixture Model (GMM) notes — the full derivation Andrew Ng works through in the unsupervised-learning block of Stanford's machine learning course. If you came here from search looking for "CS229 EM algorithm notes" or "CS229 Gaussian Mixture Model derivation," you are in the right place: every step from Jensen's inequality to the E-step responsibilities and the M-step closed-form updates is written out below.

If you want the broader overview of every CS229 lecture, see our complete CS229 lecture notes. This page is the deep dive on one specific topic: the EM algorithm and its most common application, the Gaussian Mixture Model.

Quick orientation: what EM is for and why CS229 derives it in general form first

The problem EM solves: you have data that you believe was generated by a mixture of several underlying processes, but you don't know which process generated which point. If you knew the assignments, estimating the parameters of each process would be easy (it's just maximum likelihood on a labeled subset). If you knew the parameters, figuring out the assignments would also be easy (compute which process each point is most likely to have come from). The problem is you know neither, and the two depend on each other circularly.

That's a latent variable problem: the assignment variable $z^{(i)}$ for each training example is never observed. The Expectation-Maximization algorithm is the general recipe for maximum likelihood estimation when your model has latent variables. CS229 derives it in full generality before applying it to any specific model, because the same two-step structure — guess a distribution over the latent variables, then maximize parameters against that guess — reappears constantly: in Gaussian Mixture Models, in factor analysis, in Gaussian discriminant analysis with missing labels, and in far more advanced settings like hidden Markov models and variational inference.

The problem: maximum likelihood with latent variables

Suppose you have training examples $\{x^{(1)}, \ldots, x^{(m)}\}$ and you want to fit a model $p(x, z; \theta)$ where $z$ is a latent (unobserved) variable. The log-likelihood of the observed data, marginalizing out $z$, is:

This is the object we want to maximize. The problem is immediately visible in the formula: there's a log of a sum inside the outer sum. If $z$ were observed, we'd have $\log p(x^{(i)}, z^{(i)}; \theta)$ — a log of a product of simple terms, which usually decomposes into a clean closed-form maximization. But the log of a sum over $z$ does not decompose nicely, and for most models of interest there's no closed-form maximizer of $\ell(\theta)$ directly.

EM sidesteps this by constructing, at each iteration, a tractable lower bound on $\ell(\theta)$ that touches the true log-likelihood at the current parameter estimate, then maximizing that lower bound instead. Repeat, and the log-likelihood provably never decreases.

Jensen's inequality: the tool that makes EM work

Jensen's inequality is the single piece of mathematics the whole derivation rests on. For a concave function $f$ (like $\log$) and a random variable $X$:

with equality if and only if $X$ is constant (i.e., $X = \mathbb{E}[X]$ with probability 1). Since $\log$ is concave, applying Jensen's inequality to a log-of-a-sum lets us pull the log inside, turning it into a sum-of-logs — at the cost of an inequality instead of an equality.

Here's the derivation. Let $Q_i$ be some distribution over the latent variable $z^{(i)}$ (so $Q_i(z^{(i)}) \ge 0$ and $\sum_{z^{(i)}} Q_i(z^{(i)}) = 1$). Rewrite the log-likelihood by multiplying and dividing by $Q_i(z^{(i)})$ inside the sum:

The inner sum is now an expectation over $z^{(i)} \sim Q_i$ of the quantity $\frac{p(x^{(i)}, z^{(i)}; \theta)}{Q_i(z^{(i)})}$. Applying Jensen's inequality (log is concave, so the inequality points this direction):

This right-hand side is a lower bound on $\ell(\theta)$ for any valid choice of $Q_i$. Call it $\text{ELBO}(\theta, Q)$ — the evidence lower bound. EM maximizes this lower bound instead of $\ell(\theta)$ directly, because it decomposes into a sum of logs rather than a log of a sum.

Making the bound tight: choosing $Q_i$

The bound holds for any $Q_i$, but it's only useful if it's tight at the current parameters $\theta$ — otherwise maximizing the bound doesn't guarantee progress on the real objective. Jensen's inequality is tight exactly when the random variable inside is constant, which here means:

for some constant $c$ (not depending on $z^{(i)}$). This means $Q_i(z^{(i)}) \propto p(x^{(i)}, z^{(i)}; \theta)$. Since $Q_i$ must sum to 1 over $z^{(i)}$, normalizing gives:

That's the key result: the tightest possible $Q_i$ is just the posterior distribution over the latent variable given the current parameters and the observed data. This single equation is the entire E-step, in general form, for any latent-variable model.

The EM algorithm, stated in general form

Putting the two pieces together, EM alternates:

E-step. For each $i$, set:

using the current parameters $\theta$.

M-step. Update the parameters:

Repeat until convergence. The E-step computes a tight lower bound at the current $\theta$; the M-step maximizes that bound over $\theta$, which can only move the bound (and therefore the true log-likelihood, since they agree at the old $\theta$) up.

Why EM monotonically increases the log-likelihood

This is the guarantee CS229 spends real time proving, because it's what makes EM trustworthy as an optimization procedure even though it's a heuristic rather than a closed-form solution. The argument has three steps:

1. After the E-step at iteration $t$, the bound is tight: $\text{ELBO}(\theta^{(t)}, Q^{(t)}) = \ell(\theta^{(t)})$.
2. The M-step picks $\theta^{(t+1)}$ to maximize the ELBO with $Q^{(t)}$ fixed, so $\text{ELBO}(\theta^{(t+1)}, Q^{(t)}) \ge \text{ELBO}(\theta^{(t)}, Q^{(t)})$.
3. The ELBO is always a lower bound on the true log-likelihood for any $Q$, so $\ell(\theta^{(t+1)}) \ge \text{ELBO}(\theta^{(t+1)}, Q^{(t)})$.

Chaining these:

So the log-likelihood never decreases across an EM iteration. It's not guaranteed to reach the global maximum — like k-means, EM can and does converge to a local maximum, which is why in practice you run it from several random initializations and keep the best result.

Applying EM to Gaussian Mixture Models

A Gaussian Mixture Model assumes each training example was generated by first picking a latent cluster $z^{(i)} \in \{1, \ldots, k\}$ according to a categorical distribution with parameters $\phi_j = p(z^{(i)} = j)$, and then drawing $x^{(i)}$ from the Gaussian associated with that cluster:

The parameters to estimate are the mixture weights $\phi_j$, the means $\mu_j$, and the covariances $\Sigma_j$ for $j = 1, \ldots, k$. If we knew the true cluster labels $z^{(i)}$, this would be ordinary Gaussian discriminant analysis with a closed-form maximum-likelihood solution. Because $z^{(i)}$ is latent, we apply EM.

E-step for GMM: computing the responsibilities

From the general E-step formula, $Q_i(z^{(i)} = j) = p(z^{(i)} = j \mid x^{(i)}; \phi, \mu, \Sigma)$. By Bayes' rule:

Substituting the Gaussian density for $p(x^{(i)} \mid z^{(i)} = j; \mu, \Sigma)$ and $\phi_j$ for $p(z^{(i)} = j; \phi)$:

$w_j^{(i)}$ is called the responsibility — the posterior probability that cluster $j$ generated point $x^{(i)}$, given the current parameters. Unlike k-means, which assigns each point to exactly one cluster, every point gets a soft responsibility vector across all $k$ clusters that sums to 1. This is the "soft assignment" property that distinguishes GMM/EM from k-means.

M-step for GMM: closed-form updates

With the responsibilities $w_j^{(i)}$ fixed from the E-step, the M-step maximizes $\sum_i \sum_j Q_i(z^{(i)}=j) \log \frac{p(x^{(i)}, z^{(i)}=j; \phi,\mu,\Sigma)}{Q_i(z^{(i)}=j)}$ over $\phi, \mu, \Sigma$. Because the responsibilities are now fixed numbers rather than variables, this maximization has a closed-form solution — take derivatives with respect to each parameter, set to zero, and solve (using a Lagrange multiplier for $\phi$ to enforce $\sum_j \phi_j = 1$). The results are:

Mixture weights:

The updated weight for cluster $j$ is simply the average responsibility that cluster $j$ took across all training points — the "effective number of points" assigned to $j$, divided by $m$.

Means:

This is a weighted average of all training points, where each point's contribution to cluster $j$'s mean is weighted by how responsible cluster $j$ is for it. Compare this to k-means, where each point contributes fully to exactly one centroid — here every point contributes partially to every mean.

Covariances:

The same weighted-average logic applied to the outer-product term that defines the sample covariance.

Putting it together

The full GMM training loop: initialize $\phi_j$, $\mu_j$, $\Sigma_j$ (commonly via k-means, or randomly with $k$ points from the data as initial means), then alternate:

1. E-step: compute $w_j^{(i)}$ for every point and cluster using the current parameters.
2. M-step: recompute $\phi_j$, $\mu_j$, $\Sigma_j$ using the closed-form updates above, weighted by the responsibilities just computed.

Repeat until the log-likelihood stops improving meaningfully (or you hit a fixed number of iterations). As with any EM application, run from multiple random initializations, since the log-likelihood surface for a GMM is non-convex and has many local maxima — a notoriously bad one being a component's mean collapsing exactly onto a single training point, driving that component's variance toward zero and its likelihood contribution toward infinity (degenerate solutions practical implementations guard against by adding a small regularization term to the covariance).

How EM relates to k-means

CS229 introduces k-means before EM specifically so that the comparison lands cleanly: k-means is the special case of EM-for-GMM where the E-step is forced to be a hard assignment instead of a soft one.

Concretely, if you take the GMM E-step and replace the responsibility $w_j^{(i)}$ with a one-hot vector — probability 1 on whichever cluster currently has the highest responsibility, 0 elsewhere — and additionally fix every $\Sigma_j$ to be the identity matrix (so the only thing distinguishing clusters is their mean, and "responsibility" collapses to nearest-mean-by-Euclidean-distance), the GMM M-step updates reduce exactly to the k-means update: the new mean of cluster $j$ becomes the average of the points currently hard-assigned to $j$.

This is why k-means shares EM's convergence property (monotonic decrease of the distortion objective, convergence to a local optimum only) and its practical failure mode (sensitivity to initialization). It's also why k-means is a common initializer for GMM: run k-means first to get a fast, hard-assignment estimate of cluster centers, then use those centers to initialize $\mu_j$ for the (slower, but more expressive) GMM/EM fit.

How EM connects to PCA and factor analysis, CS229's other unsupervised topics

EM, k-means, and GMMs are one part of CS229's unsupervised-learning block. The syllabus places two more topics alongside them, and it's worth being precise about how they relate rather than lumping all of unsupervised learning together:

• Principal Component Analysis (PCA) is not an EM application in the standard CS229 treatment — it has a closed-form solution via eigendecomposition of the covariance matrix (or via the singular value decomposition of the design matrix) and needs no iterative E-step/M-step procedure. PCA finds the directions of maximum variance in the data and projects onto them; CS229 derives it from both the variance-maximization view and the reconstruction-error view, and both arrive at the same eigenvectors.
• Factor analysis is the model that actually does connect back to EM. It assumes each data point is generated from a low-dimensional latent Gaussian $z$ via a linear mapping plus independent Gaussian noise: $x = \Lambda z + \mu + \epsilon$. Because $z$ is latent and continuous (not a discrete cluster label like in GMM), factor analysis is fit with a continuous-latent-variable version of EM: the E-step computes the posterior over $z^{(i)}$ given $x^{(i)}$ (which, because everything is jointly Gaussian, has a closed form), and the M-step updates $\Lambda$, $\mu$, and the noise covariance.

So the throughline for CS229's unsupervised block is: k-means (hard, no probabilistic model) generalizes to GMM (soft, EM with discrete latent variable), and GMM generalizes conceptually to factor analysis (EM with continuous latent variable), while PCA sits to the side as a separate, closed-form dimensionality-reduction technique that doesn't require EM at all. If you're trying to keep the four topics straight for an exam: ask whether the latent variable is a cluster label (GMM) or a continuous vector (factor analysis), and whether the fitting procedure is iterative (EM, for both GMM and factor analysis) or a single eigendecomposition (PCA).

Common derivation pitfalls

Three places where the EM derivation trips people up on CS229 problem sets:

1. Forgetting why Jensen's inequality points the direction it does. $\log$ is concave, so $\log(\mathbb{E}[X]) \ge \mathbb{E}[\log X]$ — the inequality gives a lower bound. If you mix this up you'll conclude EM maximizes an upper bound, which would break the monotonic-improvement guarantee entirely.

2. Treating the E-step as an approximation rather than an exact posterior. The E-step isn't a heuristic guess at $Q_i$ — it's the exact, provably tightest choice, derived directly from the equality condition of Jensen's inequality. It's exact given the current parameters; it just doesn't stay exact once the M-step moves the parameters (which is exactly why you have to redo the E-step every iteration).

3. Assuming EM finds the global maximum. It doesn't, and the derivation doesn't claim it does — only that each iteration doesn't decrease the log-likelihood. For a genuinely non-convex log-likelihood surface (which GMMs have), different initializations can converge to meaningfully different local optima. This is why running EM multiple times with different seeds and keeping the highest final log-likelihood is standard practice, not a workaround for a bug.

What is the EM algorithm used for in CS229?

EM is CS229's general tool for maximum likelihood estimation when the model has a latent (unobserved) variable — most concretely, fitting Gaussian Mixture Models, but the same E-step/M-step structure also applies to factor analysis and to mixtures of other distributions. It alternates between computing a posterior distribution over the latent variable given the current parameters (E-step) and maximizing the parameters against that posterior (M-step).

Why does EM use Jensen's inequality?

Directly maximizing the log-likelihood with a latent variable requires optimizing a log of a sum, which generally has no closed form. Jensen's inequality (applied because $\log$ is concave) lets you construct a lower bound that turns the log of a sum into a sum of logs — a tractable quantity — and shows that choosing the latent-variable posterior as $Q_i$ makes that bound exactly tight at the current parameters. Maximizing the tight bound then provably does not decrease the true log-likelihood.

What's the difference between the E-step and the M-step?

The E-step computes, for each training example, the posterior probability distribution over its latent variable given the current model parameters (for a GMM, this is the "responsibility" of each cluster for that point). The M-step then treats those posteriors as fixed and re-solves for the model parameters that maximize the expected complete-data log-likelihood — for a GMM this has closed-form updates for the mixture weights, means, and covariances.

How is k-means related to the EM algorithm?

K-means is EM's hard-assignment special case. Where GMM's E-step computes a soft responsibility for every cluster, k-means forces a one-hot (hard) assignment to the single nearest centroid, and where GMM tracks a full covariance per cluster, k-means implicitly fixes every cluster's covariance to the identity matrix. With those two restrictions, the GMM M-step's weighted-mean update collapses exactly into the k-means centroid update.

Does EM always converge to the best solution?

No. EM guarantees the log-likelihood never decreases from one iteration to the next, and it does converge to a stationary point, but that point can be a local maximum rather than the global one — the same limitation k-means has. For Gaussian Mixture Models, the log-likelihood surface is non-convex, so different random initializations can land on different local optima. Standard practice is to run EM several times from different initializations (often seeded with a k-means run) and keep the fit with the highest final log-likelihood.

Turn any Stanford lecture into notes like these

These notes were synthesized from the CS229 unsupervised-learning lectures. Notiq does the same thing for any YouTube lecture — paste a link, get back structured notes with the derivations written out, flashcards, exam questions, and worked examples. Try the Stanford CS229 lecture library in Notiq or paste your own lecture.