Gaussian Mixture Model (EM) - Unsupervised Learning

Implement Gaussian Mixture Model (EM) in Python with NumPy. hard exercise with examples, hints, and instant feedback. No setup or signup required.

Gaussian Mixture Model (EM)

What you'll build

Two functions for a 1-D Gaussian mixture: a numerically stable E-step that computes responsibilities in log space, and the full EM loop that fits the weights, means and variances from given initial values.

Background

A Gaussian mixture models the density of x as a weighted sum of K Gaussians. It is "soft k-means": instead of a hard label, each point gets a probability (responsibility) for each component. EM alternates between computing responsibilities from the current parameters (E-step) and re-estimating the parameters from the responsibilities (M-step), and the log-likelihood never decreases. scikit-learn's GaussianMixture does the same with full covariance matrices in higher dimensions.

Model

p(x) = sum_k  weights[k] * N(x | means[k], variances[k]),      sum_k weights[k] = 1
log N(x | mu, var) = -0.5 * log(2*pi*var) - (x - mu)^2 / (2*var)

1. `gmm_e_step(x, means, variances, weights)`

log_joint[i, k] = log(weights[k]) + log N(x[i] | means[k], variances[k])     # (n, K)
log_px[i]       = logsumexp_k log_joint[i, k]                               # log p(x_i)
resp[i, k]      = exp(log_joint[i, k] - log_px[i])
log_likelihood  = sum_i log_px[i]
logsumexp(a)    = max(a) + log(sum(exp(a - max(a))))

2. `gmm_em(x, means, variances, weights, n_iters, reg_covar=1e-6)`

Repeat n_iters times:

E-step at the current parameters: resp, ll = gmm_e_step(...); append ll to the history.
M-step, with N_k = sum_i resp[i, k]:

weights[k]   = N_k / n
means[k]     = sum_i resp[i, k] * x[i] / N_k
variances[k] = sum_i resp[i, k] * (x[i] - means[k])^2 / N_k  +  reg_covar     # uses the NEW means

Inputs

x: 1-D float array, shape (n,)
means, variances, weights: initial parameters, each shape (K,); variances > 0, weights sum to 1
n_iters: number of EM iterations (>= 1); reg_covar: small float added to every variance

Expected Return Format

gmm_e_step returns a tuple (resp, log_likelihood): resp has shape (n, K) with rows summing to 1, and log_likelihood is a Python float.
gmm_em returns a dictionary with these keys:
- 'means', 'variances', 'weights': final parameters, each shape (K,)
- 'log_likelihoods': list of n_iters floats; entry t is the log-likelihood before the t-th M-step (entry 0 belongs to the initial parameters)

Edge cases and conventions

Underflow: for a point far from every component all densities are 0.0 in floating point, and the naive pdf / pdf.sum() gives 0/0 = nan. Log space with logsumexp avoids it.
reg_covar (default 1e-6, as in scikit-learn) stops a component from collapsing onto a single point with variance 0.
The log-likelihood history is non-decreasing (up to floating-point error).

Constraints

NumPy only, vectorised over points and components.

Follow-up

How would you choose K? (BIC/AIC, held-out likelihood.)
Show that k-means is the limit of EM with equal, shared variances shrinking to 0.

Examples

Input: gmm_e_step(np.array([0.0, 4.0]), np.array([0.0, 4.0]), np.array([1.0, 1.0]), np.array([0.5, 0.5]))
Output: (array([[0.999665, 0.000335], [0.000335, 0.999665]]), -3.223501)

Each point is almost entirely explained by the component centred on it (values rounded to 6 decimals).

Input: gmm_em(x, np.array([-1.0, 1.0]), np.array([1.0, 1.0]), np.array([0.5, 0.5]), 30), where x has 60 points around -3 and 40 points around 3
Output: means ≈ [-2.9, 3.0], weights ≈ [0.6, 0.4]

Starting near the middle, EM pulls each component onto one cluster and the weights converge to the cluster fractions.

Practice in Python

import numpy as np

def gmm_e_step(x, means, variances, weights):
    """
    E-step for a 1-D Gaussian mixture, computed in log space.

    Args:
        x: Data, shape (n,)
        means, variances, weights: Component parameters, each shape (K,)

    Returns:
        (resp, log_likelihood): responsibilities (n, K) and float sum_i log p(x_i)
    """
    # Your code here
    pass


def gmm_em(x, means, variances, weights, n_iters=10, reg_covar=1e-6):
    """
    Fit a 1-D Gaussian mixture with EM from the given initial parameters.

    Args:
        x: Data, shape (n,)
        means, variances, weights: Initial parameters, each shape (K,)
        n_iters: Number of EM iterations
        reg_covar: Added to every variance in the M-step

    Returns:
        dict with 'means', 'variances', 'weights' (each (K,)) and
        'log_likelihoods' (list of n_iters floats, before each M-step)
    """
    # Your code here
    pass

The interactive editor loads with Python and NumPy. Run your implementation against 6 tests, or use the progressive hints.

More exercises in this section →

Gaussian Mixture Model (EM) - Unsupervised Learning