Expectation-Maximization for Gaussian Mixtures: The Full Math and Code

Imagine a pile of measurements, and every one of them came from one of several groups, but nobody wrote down which. You must learn the shape of each group without ever seeing a label. It sounds like a chicken and egg problem, and it is. Today we solve it with one of the most beautiful algorithms in machine learning.

• Data with hidden group labels
• Can we still learn the groups?

In the last lesson we met the support vector machine dual and the kernel trick. This lesson returns to the Gaussian mixture model that we only sketched in lesson 22. There we used it as a tool. Here we open the box and derive exactly why the algorithm works, with numbers, proofs and code.

• Lesson 43: SVM dual
• Lesson 44: EM for GMMs

Here is the plan. First the hidden label problem and why plain maximum likelihood struggles. Then responsibilities, the E-step and the M-step, with a full example by hand. Next we derive the update formulas, sketch the proof that the likelihood never goes down, write the code in NumPy, and finish with pitfalls and k-means as a special case.

• Hidden labels and responsibilities
• E-step, M-step, derivation
• ELBO, code, choosing K

A Gaussian mixture tells a two step story. First, nature secretly picks a component number k, with probability pi sub k. Then it draws the point x from the Gaussian with mean mu sub k and covariance Sigma sub k. We see only x. The choice of k is called a latent variable, written z, because it is real but hidden.

• Pick a component k with probability pi_k
• Draw x from Gaussian k

Summing over the hidden choice gives the density of x. It is p of x equals the sum over k of pi sub k times a Gaussian density with parameters mu sub k and Sigma sub k. The weights pi are non negative and add to one. Think of it as a weighted blend of bell curves, giving one bumpy curve with several peaks.

• p(x) = sum of pi_k times N(x | mu_k, Sigma_k)
• Weights are positive and sum to one

Maximum likelihood says: maximise the log of the probability of all the data. With independent points that is the sum over n of the log of a sum over k. Here is the trouble. A log of a product splits nicely, but a log of a sum does not. If we differentiate, the inner sum stays glued together, and we get no clean closed form.

• Log-likelihood: sum_n log sum_k ...
• The log cannot enter the inner sum

Notice what would happen if the labels were known. We would simply group the points, then take each group's mean, covariance and share of the data. Easy. Conversely, if we knew the parameters, we could guess each point's group. Labels need parameters, and parameters need labels. EM breaks the circle by taking turns.

• Known labels give easy estimates
• Hidden labels give a circular problem

Given current parameters, how likely is it that point x sub n came from component k? Bayes' theorem from lesson 8 answers. The prior is pi sub k, the likelihood is the Gaussian density at x sub n, and we divide by the sum of the same quantity over all components. The result is called the responsibility, r sub n k.

• r_nk = P(z = k | x_n)
• Bayes: prior times likelihood, normalised

Responsibilities are soft. A point between two clusters might be seventy percent component one and thirty percent component two. For every point, the responsibilities across components add up to one. A common mistake is to round them to zero or one too early. Keeping them soft is what makes the math smooth, and it carries uncertainty forward.

• Each row of r sums to one
• A point can belong 70/30

The algorithm is a loop. In the E-step, the expectation step, we hold the parameters fixed and compute all responsibilities. In the M-step, the maximisation step, we hold the responsibilities fixed and re-fit the means, covariances and weights as if those soft labels were real. Repeat until the parameters stop changing.

• E-step: compute responsibilities
• M-step: re-fit parameters
• Repeat until stable

Let us run one full iteration by hand in one dimension. Our four data points are zero, one, three and four. Start with a rough guess: component one has mean zero, component two has mean four, both have variance one, and both weights are one half. Because variances and weights match, the shared constants cancel and we only need the exponential part.

• Data: 0, 1, 3, 4
• mu = 0 and 4, variance 1, pi = 0.5 each

Take x equals one. Distance to mean zero is one, so the exponent is minus one squared over two, giving e to the minus point five, which is about point six oh six five. Distance to mean four is three, so e to the minus four point five, about point oh one one one. Dividing point six oh six five by their sum, point six one seven six, gives a responsibility of point nine eight two for component one.

• Comp 1: exp(-0.5) = 0.6065
• Comp 2: exp(-4.5) = 0.0111
• r_1 = 0.982

By the same steps, the point zero gets responsibility point nine nine nine seven for component one. The point three mirrors the point one, so component one gets only point oh one eight. The point four is almost entirely component two. Notice how fast the odds fall off: a distance of three standard deviations is already nearly decisive.

• x=0: 0.9997 / 0.0003
• x=1: 0.982 / 0.018
• x=3: 0.018 / 0.982
• x=4: 0.0003 / 0.9997

Now the M-step. The effective number of points in component one, called N sub one, is the sum of its responsibilities, which is two. The new mean is the responsibility weighted average of the data. That is point nine eight two times one, plus point oh one eight times three, plus point oh oh oh three times four, all over two. This gives about point five one nine. By symmetry, mu two becomes about three point four eight one.

• N_1 = 0.9997+0.982+0.018+0.0003 = 2.0
• mu_1 = weighted average
• mu_1 = 0.519

For the variance we take responsibility weighted squared distances from the new mean. For component one these are point two six eight eight, plus point two two seven six, plus point one one oh eight, plus point oh oh three six, which sums to point six one one. Divide by N sub one, which is two, and the variance is about point three oh five. The weight is N sub one over N, two over four, so one half.

• Weighted squared distances
• sigma_1^2 = 0.305
• pi_1 = N_1 / N = 2/4 = 0.5

One pass already moved the means inward and shrank the variances from one to about point three. With the tighter curves, the next E-step is even more confident, so the clusters tighten again. After a few more rounds the changes become tiny and we stop. A common mistake is to check only the parameters. Many people instead watch the log likelihood, which we will soon see can only go up.

• Variance dropped from 1.0 to 0.31
• Another round tightens further
• Stop when change is tiny

Now let us derive those update rules properly. Pretend the labels z were known. Then each point contributes only its own component, and the log of a single product splits into a sum: log of pi sub k plus log of the Gaussian density. That is the complete data log likelihood, and it is easy to maximise. The only problem is that z is hidden.

• If z were known, the log splits
• Sum of ln pi_k + ln N(...)

Since we do not know z, we take its expectation under the responsibilities from the E-step. The hard indicator that point n belongs to component k is replaced by the soft number r sub n k. The objective, called Q, becomes the double sum over n and k of r sub n k times the bracket of log pi sub k plus log of the Gaussian. The responsibilities are now constants, so this is easy to differentiate.

• Replace indicators by r_nk
• Q = sum_n sum_k r_nk [ln pi_k + ln N]

Differentiate Q with respect to mu sub k. Only the Gaussian log term depends on it, and its gradient is the inverse covariance times x minus mu, a result that follows from the multivariate Gaussian in lesson 38. Set the sum to zero, and the inverse covariance factors out and cancels. What remains says the responsibility weighted sum of x minus mu is zero, so mu sub k is the weighted average, exactly as we computed by hand.

• Gradient of ln N wrt mu: Sigma^-1 (x - mu)
• Set sum_n r_nk (x_n - mu_k) = 0
• mu_k = (1/N_k) sum_n r_nk x_n

For the covariance, the log density contains minus one half log determinant of Sigma, and minus one half the quadratic form. It is easiest to differentiate with respect to the inverse of Sigma. The determinant term gives a half of Sigma, the quadratic term gives minus a half of the outer product x minus mu times itself transposed. Setting the sum to zero gives Sigma sub k equal to the responsibility weighted average of those outer products.

• Differentiate wrt Sigma^-1
• Result: weighted scatter matrix
• Sigma_k = (1/N_k) sum r_nk (x-mu)(x-mu)^T

The weights pi have a constraint: they must sum to one. This is a job for the Lagrange multipliers of lesson 37. We maximise the sum of N sub k times log pi sub k plus lambda times the constraint. Differentiating gives N sub k over pi sub k plus lambda equals zero. So pi sub k is proportional to N sub k, and the constraint forces lambda to equal minus N. The answer is pi sub k equals N sub k over N.

• Maximise sum N_k ln pi_k
• Constraint: sum pi_k = 1
• Lagrange gives pi_k = N_k / N

So the M-step is just weighted versions of the familiar formulas. Effective count, weighted mean, weighted covariance, and share of the total. A frequent mistake is to use the new mean when computing the covariance. Actually that is correct and required: compute the mean first, then the covariance around it. Another mistake is dividing by N instead of N sub k, which would shrink every covariance wrongly.

• Effective count N_k
• Weighted mean, weighted covariance
• Weights are fractions of total

Why does this loop always make progress? Pick any distribution q over the hidden labels. The true log likelihood splits exactly into two parts: a lower bound called the ELBO, short for evidence lower bound, plus a KL divergence between q and the true posterior over labels. Since a KL divergence is never negative, the ELBO can never exceed the log likelihood.

• ln p(X) = ELBO(q) + KL(q || posterior)
• KL is never negative

Where does the bound come from? Multiply and divide inside the log by q, so the sum becomes an average under q. The log is concave, so its value at an average is at least the average of its values. That is Jensen's inequality. Pushing the log inside the sum gives exactly the ELBO, and the gap is the KL divergence.

• Log is concave
• ln E[f] >= E[ln f]
• This produces the lower bound

Now the two steps get a clear meaning. The E-step chooses q to be exactly the posterior, which are our responsibilities. That makes the KL term zero, so the ELBO touches the true log likelihood at the current parameters. We have built a tight lower bound, a curve that sits underneath the true one and kisses it at one point.

• Choose q = posterior p(z | x)
• KL becomes zero
• Bound touches the likelihood

The M-step maximises the bound over the parameters. Moving to the top of the green curve raises the bound, and since the true likelihood is at least the bound, the true likelihood at the new point is at least as high as where we started. So each full iteration can only keep or increase the log likelihood. That is the proof sketch, and it is why monitoring the log likelihood is such a good debugging tool.

• Maximise the ELBO over theta
• New point is higher on the bound
• True likelihood is higher still

Be careful about what this guarantee does not say. It promises the likelihood never decreases, and that it converges to a stationary point. It does not promise the best possible solution. The landscape has many bumps, and EM climbs whichever hill it starts on. That is why initialization matters so much, as we will see in a moment.

• Monotone increase, yes
• Global optimum, no
• Depends on initialization

Time to code it. We store the data as an N by d array. For initialisation, pick K random data points as the means, set every covariance to the identity matrix, and give every component the weight one over K. Then we build a loop that alternates the two steps. Using real data points as starting means keeps the initial components near actual data.

• Random data points as initial means
• Identity covariances, equal weights

In the E-step, compute the log of pi sub k times the Gaussian density for every point and component, giving an N by K array. Do not exponentiate directly. Densities can underflow to zero in high dimensions, and then we divide zero by zero. Instead subtract each row's maximum, exponentiate, and normalise. This is the log sum exp trick, and it also gives the log likelihood for free.

• Work with log densities
• Subtract the row maximum
• Avoids underflow to zero

The M-step is mostly matrix multiplication. The effective counts are the column sums of r. The means are r transposed times X, divided by the counts. For each component, subtract its mean from the data, then the covariance is the weighted sum of outer products, which is one line using the transpose times a weighted copy of the centred data. Finally the weights are the counts over N. Add a tiny value to the diagonal for stability.

• Nk = r.sum(axis=0)
• mu = r.T @ X / Nk
• Sigma from weighted outer products

Wrap both steps in a loop. After every iteration, record the log likelihood. Stop when the improvement drops below a tolerance such as ten to the minus six, or when you hit a maximum number of iterations. A great unit test: assert that the log likelihood never decreases, apart from tiny floating point noise. If it does decrease, you have a bug, usually in the covariance update.

• Repeat E and M steps
• Track log-likelihood each iteration
• Stop when improvement < tolerance

Now compare with scikit-learn's GaussianMixture, using the same initial means and covariance type full. After fitting, the means, covariances and weights should match ours up to a permutation of the component labels, since component order is arbitrary. The final log likelihood should agree to several decimals. The library adds extras, such as multiple restarts, smarter initialisation with k-means, and a regularisation term on the covariance.

• GaussianMixture(n_components=K)
• Match means up to label order
• Likelihoods agree

Run our code with ten different seeds and you may get several different answers. One seed may place two components on a single cluster and leave another cluster covered by a broad component. The cure is simple. Run multiple restarts, compare their final log likelihoods, and keep the best. Initialising with k-means, which sklearn does by default, also helps a lot.

• Different seeds, different answers
• Run several restarts
• Keep the best log-likelihood

Here is the nastiest failure. Suppose one component's mean lands exactly on a single data point. If its variance shrinks toward zero, the density at that point grows without limit, and the likelihood shoots to infinity. This is a singularity, and it is a bad solution, not a good one. The standard fix is to add a small constant to the covariance diagonal, or to restart any component that collapses.

• A component collapses onto one point
• Variance goes to zero, density to infinity
• Fix: add reg_covar to the diagonal

How many components should we use? More components always raise the likelihood, so likelihood alone would pick a huge K. Information criteria add a penalty. The Akaike criterion is two p minus two log L, where p is the number of free parameters and L is the maximised likelihood. The Bayesian criterion replaces two p with p times log N, so it penalises complexity more once N is above about eight. Choose the K with the lowest score.

• BIC = p ln N - 2 ln L
• AIC = 2p - 2 ln L
• Lower is better

Counting parameters is a place where people slip. In two dimensions with three components and full covariances, each mean has two numbers, so six in total. Each symmetric two by two covariance has three free numbers, so nine. The weights add to one, so only two are free. The total is seventeen. Switching to diagonal covariances cuts this to eleven, a trade between flexibility and parameter count.

• d=2, K=3, full covariance
• Means: 6, covariances: 9, weights: 2
• p = 17

Now the connection promised earlier. Force every component to have the same spherical covariance, sigma squared times the identity, and equal weights. Then shrink sigma squared toward zero. In our example, with variance point two five, the responsibility of point one for component two is around one in ten million. Soft assignments snap to hard ones: each point goes fully to its nearest mean.

• Sigma_k = sigma^2 I, shared and equal
• Let sigma^2 go to zero
• Responsibilities become 0 or 1

In that limit, the E-step becomes the k-means assignment step, nearest mean wins, and the M-step becomes the update step, each mean is the average of its points. So k-means is EM for a very restricted mixture. A full Gaussian mixture adds three things: soft memberships, separate cluster sizes, and elliptical shapes. That is why it handles stretched clusters where k-means fails.

• E-step becomes nearest-mean assignment
• M-step becomes cluster averages
• GMM adds shape, size and softness

Let us recap. Hidden labels turn the log likelihood into a log of a sum with no closed form. Responsibilities come from Bayes' theorem. The E-step computes them, and the M-step re-fits means, covariances and weights with responsibility weights. The ELBO and Jensen's inequality show the likelihood never decreases. In practice, use log sum exp, restarts, regularised covariances, and BIC to pick K. And k-means is the hard limit.

• Hidden labels, log of a sum
• Responsibilities, E-step, M-step
• ELBO guarantees monotone climb

In the next lesson, number forty five, we turn to choosing the settings that learning algorithms cannot learn for themselves. Grid search, random search, and Bayesian optimization. See you in lesson forty five.

• Lesson 45: tuning with search
• Grid, random, Bayesian optimization
• Subscribe for the full course