ML coding · Probabilistic models
One EM step for a Gaussian mixture
A 1-D Gaussian mixture has weights pi, means mu and variances
var (each of length K). Given data x, perform one EM iteration and return
[new_pi, new_mu, new_var, loglik]:
- E-step: rik = πkN(xi; μk, σ²k) / Σj πjN(xi; μj, σ²j).
- M-step: Nk = Σi rik; new πk = Nk/n; new μk = Σi rikxi/Nk; new σ²k = Σi rik(xi − new μk)²/Nk.
- loglik: Σi log Σk πkN(xi; μk, σ²k), under the parameters you were given.
Examples
Input x = [-1, 1], pi = [0.5, 0.5], mu = [0, 0], var = [1, 1]
Output [[0.5, 0.5], [0, 0], [1, 1], -2.8379]
Two identical components: each takes half of every point, and the variance becomes the data's.
Input x = [0, 0.2, 10, 10.4], pi = [0.5, 0.5], mu = [0, 10], var = [1, 1]
Output [[0.5, 0.5], [0.1, 10.2], [0.01, 0.04], -6.5483]
Well separated: each component takes its own cluster.
Input x = [0, 1, 2, 9, 10, 11], pi = [0.3, 0.7], mu = [1, 9], var = [2, 2]
Output [[0.5, 0.5], [1, 10], [0.6667, 0.667], -14.025]
The weights move towards the true half-and-half split.
Constraints
- 1 ≤ n ≤ 10⁵, 1 ≤ K ≤ 10
- All weights and variances positive; no component receives zero total responsibility
- Checked to a relative tolerance of 10⁻⁹
Hints
Hint 1
Build the (n, K) matrix of log π_k + log N(x_i; μ_k, σ²_k). Broadcasting x[:, None] against mu does it.
Hint 2
Normalise each row with log-sum-exp: r = exp(logp − lse). The sum over rows of lse is the log-likelihood.
Hint 3
The M-step is three weighted averages, with the columns of r as weights.
Where candidates go wrong
- Computing the densities directly and dividing: a point far from every component gives 0/0.
- Using the old means in the variance update. The new variance is around the new mean.
- Reporting the log-likelihood after the update rather than under the parameters given.
- Dividing by n instead of N_k in the mean and variance.
The idea
The E-step is Bayes' rule for each point, done in log space: the log of πk times a normal density is log πk − ½ log(2πσ²k) − (x − μk)²/(2σ²k), and the row-wise log-sum-exp normalises it without ever forming a probability that could underflow. The same log-sum-exp, summed over points, is the log-likelihood.
The M-step maximises the expected complete-data log-likelihood, which for Gaussians gives weighted versions of the usual estimates. EM guarantees the likelihood of the new parameters is at least that of the old ones.
Time O(nK) · Space O(nK)
Reference solution
import numpy as np
def em_step(x, pi, mu, var):
x = np.asarray(x, dtype=float)[:, None] # (n, 1)
pi, mu, var = (np.asarray(a, dtype=float) for a in (pi, mu, var))
logp = np.log(pi) - 0.5 * np.log(2 * np.pi * var) - (x - mu) ** 2 / (2 * var) # (n, K)
m = logp.max(axis=1, keepdims=True)
lse = m + np.log(np.exp(logp - m).sum(axis=1, keepdims=True)) # log p(x_i)
r = np.exp(logp - lse) # responsibilities
nk = r.sum(axis=0)
new_mu = (r * x).sum(axis=0) / nk
new_var = (r * (x - new_mu) ** 2).sum(axis=0) / nk # around the new means
return [nk / len(x), new_mu, new_var, float(lse.sum())]Follow-ups an interviewer asks
- Run EM to convergence and stop when the log-likelihood improves by less than 10⁻⁸.
- Extend to d dimensions with full covariance matrices.
- Add a variance floor and explain the degenerate solution it prevents.
- Compare with k-means: show that k-means is EM with equal, vanishing variances and hard assignments.
Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.
Why interviewers ask it
EM is the standard example of latent-variable fitting, and research interviewers use it to see whether a candidate can derive an update rule, not only apply one. Writing a single step cleanly, with the responsibilities in log space, is the core of it.
The theory behind it is in the lesson Naive Bayes, EM, Viterbi and Markov chains from scratch.
Keep going
Gaussian naive Bayes classifier
Fit class priors, means and variances, then predict by the largest log posterior, vectorised over test points and classes.
Open the problem → Medium20 minOne EM step for a Gaussian mixture
Compute responsibilities in log space, then refit weights, means and variances: one expectation-maximisation step for a 1-D Gaussian mixture, from scratch..
Open the problem → Medium25 minThe Viterbi algorithm for a hidden Markov model
Find the most likely sequence of hidden states in an HMM with dynamic programming in log space, with back-pointers to recover the path.
Open the problem → Easy12 minStationary distribution of a Markov chain
Solve πP = π with Σπ = 1 by replacing one redundant equation, in NumPy.
Open the problem →