Open In Colab

Introduction

Expectation Maximization (EM) is one of the most elegant optimization algorithms in machine learning.

It appears everywhere:

  • Gaussian Mixture Models (GMMs)

  • Hidden Markov Models

  • Latent Class Analysis

  • Missing data estimation

  • Topic models

  • Variational inference foundations

In this tutorial, we will:

  1. Build EM for Gaussian Mixture Models from scratch

  2. Derive the mathematics carefully

  3. Implement every step manually in NumPy

  4. Reinforce everything using scikit-learn

  5. Analyze likelihoods, AIC, BIC, responsibilities, and latent profiles

  6. Visualize hidden structure in physiological sleep data

The tutorial uses WHOOP biometric sleep data (Yoo et al., 2021) and demonstrates real probabilistic clustering.

Why EM Exists

Suppose your data comes from multiple hidden populations.

For example:

  • athletes recovering well

  • stressed individuals

  • sleep-deprived individuals

But you do not know which row belongs to which group.

Instead of hard labels, we assume:

\[z_i \in {1,2,\dots,K}\]

where:

  • \(z_i\) is a hidden latent variable

  • \(K\) is the number of hidden clusters

We only observe:

\[x_i\]

The hidden labels are missing.

That is precisely where EM shines.

Gaussian Mixture Models

A Gaussian Mixture Model assumes data is generated from:

\[p(x_i) = \sum_{k=1}^{K} \pi_k \mathcal{N}(x_i \mid \mu_k, \Sigma_k)\]

Where:

  • \(\pi_k\) = mixing weight

  • \(\mu_k\) = cluster mean

  • \(\Sigma_k\) = covariance matrix

The constraints are:

\[\sum_{k=1}^{K} \pi_k = 1\]

and

\[0 \le \pi_k \le 1\]

The Gaussian Density

The multivariate Gaussian density is:

\[\mathcal{N}(x\mid\mu,\Sigma)=\frac{1}{\sqrt{(2\pi)^d|\Sigma|}}\exp\left(-\frac12(x-\mu)^T\Sigma^{-1}(x-\mu)\right)\]

This equation defines the probability density of a point under a Gaussian distribution.

Key pieces:

  • determinant $$ \Sigma $$

  • inverse covariance \(\Sigma^{-1}\)

  • Mahalanobis distance

Dataset

We use biometric sleep features:

features = [
    "rhr",
    "hrv",
    "hours.of.sleep",
    "sleep.disturbances",
    "latency.min",
]

These represent:

  • resting heart rate

  • heart rate variability

  • total sleep duration

  • disturbances

  • sleep latency

Loading the Data

import pandas as pd

url = "https://raw.githubusercontent.com/OpenNeuroDatasets/ds003823/main/phenotype/Whoop_Summary.tsv"

df = pd.read_csv(url, sep="\t")

display(df.head())

Preparing the Matrix

We define:

\[X \in \mathbb{R}^{n \times d}\]

where:

  • \(n\) is the number of observations

  • \(d\) is the number of features

import numpy as np

X = df[
    [
        "rhr",
        "hrv",
        "hours.of.sleep",
        "sleep.disturbances",
        "latency.min"
    ]
].dropna().to_numpy()

Why Standardization Matters

Gaussian mixtures are sensitive to scale.

Without normalization:

  • HRV may dominate

  • sleep latency may become irrelevant

  • covariance estimation becomes unstable

We standardize:

\[x_{scaled} = \frac{x - \mu}{\sigma}\] 
X_scaled = (X - X.mean(axis=0)) / X.std(axis=0)

EM Step 1 – Initialization

We initialize:

  • means

  • covariance matrices

  • mixing weights

K = 3

n, d = X_scaled.shape

random_rows = np.random.choice(n, K, replace=False)

mu = X_scaled[random_rows]

Sigma = np.array([np.eye(d) for _ in range(K)])

pi = np.ones(K) / K

Initially:

\[\Sigma_k = I\]

and

\[\pi_k = \frac1K\]

Building the Gaussian PDF

def gaussian_pdf(x, mu, Sigma):

    d = len(mu)

    Sigma = Sigma + 1e-6 * np.eye(d)

    det = np.linalg.det(Sigma)
    inv = np.linalg.inv(Sigma)

    diff = x - mu

    return (
        1 / np.sqrt((2 * np.pi) ** d * det)
        * np.exp(-0.5 * diff.T @ inv @ diff)
    )

Notice the stabilization trick:

Sigma + 1e-6 * np.eye(d)

This prevents singular covariance matrices.

The E-Step

The E-step computes responsibilities:

\[r_{ik} = P(z_i = k \mid x_i)\]

Using Bayes theorem:

\[r_{ik}=\frac{\pi_k\mathcal{N}(x_i\mid\mu_k,\Sigma_k)}{\sum_{j=1}^K\pi_j\mathcal{N}(x_i\mid\mu_j,\Sigma_j)}\]

Responsibilities are soft cluster assignments.

E-Step in Code

r = np.zeros((n, K))

for i in range(n):

    for k in range(K):

        r[i, k] = (
            pi[k]
            * gaussian_pdf(X_scaled[i], mu[k], Sigma[k])
        )

    r[i, :] = r[i, :] / r[i, :].sum()

Each row sums to 1:

\[\sum_k r_{ik} = 1\]

The M-Step

Now we update parameters using weighted statistics.

Updating Means

\[\mu_k = \frac{ \sum_i r_{ik}x_i }{ \sum_i r_{ik} }\] 
mu[k] = np.sum(
    r[:, k].reshape(-1, 1) * X_scaled,
    axis=0
) / N_k[k]

Updating Covariances

\[\Sigma_k = \frac{ \sum_i r_{ik}(x_i-\mu_k)(x_i-\mu_k)^T }{ N_k }\] 
Sigma[k] = np.zeros((d, d))

for i in range(n):

    diff = (X_scaled[i] - mu[k]).reshape(-1, 1)

    Sigma[k] += r[i, k] * (diff @ diff.T)

Sigma[k] = Sigma[k] / N_k[k]

Updating Mixing Weights

\[\pi_k = \frac{N_k}{n}\]

where

\[N_k = \sum_i r_{ik}\] 
pi[k] = N_k[k] / n

Log Likelihood

EM maximizes the log likelihood:

\[\log L=\sum_{i=1}^n\log\left(\sum_{k=1}^K\pi_k\mathcal{N}(x_i\mid\mu_k,\Sigma_k)\right)\]

Why EM Monotonically Improves Likelihood

EM maximizes a lower bound on the log-likelihood using latent-variable expectations.

The complete-data likelihood is:

\[\log p(X, Z \mid \theta)\]

Since \(Z\) is unobserved, EM optimizes the expected complete-data log likelihood:

\[Q(\theta \mid \theta^{old}) \mathbb{E}_{Z \mid X,\theta^{old}} \left[ \log p(X,Z \mid \theta) \right]\]

The E-step computes the expectation.

The M-step maximizes this expectation with respect to \(\theta\).

This guarantees non-decreasing likelihood at every iteration.

Computing It

log_likelihood = 0

for i in range(n):

    total = 0

    for k in range(K):

        total += (
            pi[k]
            * gaussian_pdf(X_scaled[i], mu[k], Sigma[k])
        )

    log_likelihood += np.log(total)

Full EM Loop

for iteration in range(max_iter):

    # E-step
    ...

    # M-step
    ...

    # log likelihood
    ...

    if abs(log_likelihood - old_log_likelihood) < tol:
        break

EM alternates between:

  • estimating hidden variables

  • optimizing parameters

until convergence.

Final Cluster Assignments

clusters = np.argmax(r, axis=1)

This converts soft assignments into hard labels.

Geometric Interpretation

K-Means partitions space using hard Voronoi boundaries.

Gaussian Mixture Models instead learn probability densities with full covariance structure.

This allows:

  • elliptical clusters

  • rotated clusters

  • anisotropic variance

  • uncertainty estimation

The covariance matrix determines:

  • cluster orientation

  • cluster spread

  • feature correlation

Why Scikit-Learn Matters

Writing EM manually teaches the mathematics.

But production systems need:

  • numerical stability

  • vectorization

  • optimized covariance estimation

  • convergence diagnostics

  • multiple random restarts

That is where scikit-learn GaussianMixture becomes invaluable.

Using Scikit-Learn GaussianMixture

from sklearn.mixture import GaussianMixture

gmm = GaussianMixture(
    n_components=3,
    covariance_type="full",
    random_state=42,
    n_init=20
)

gmm.fit(X_scaled)

This internally performs EM.

Reinforcing the Mathematics

Scikit-learn exposes the exact quantities we derived manually.

Means

gmm.means_

Corresponds to:

\[\mu_k\]

Covariances

gmm.covariances_

Corresponds to:

\[\Sigma_k\]

Mixing Weights

gmm.weights_

Corresponds to:

\[\pi_k\]

Responsibilities

gmm.predict_proba(X_scaled)

Corresponds to:

\[r_{ik}\]

Choosing the Number of Clusters

A huge issue in latent-variable modeling is selecting \(K\).

Too small:

  • underfitting

Too large:

  • overfitting

We use:

  • AIC

  • BIC

Akaike Information Criterion

\[AIC = 2p - 2\log L\]

Where:

  • \(p\) = number of parameters

  • \(L\) = likelihood

Bayesian Information Criterion

\[BIC = p\log(n) - 2\log L\]

BIC penalizes complexity more aggressively.

Testing Different Numbers of Clusters

results = []

for k in range(1,10):

    gmm = GaussianMixture(
        n_components=k,
        covariance_type="full",
        random_state=42,
        n_init=20
    )

    gmm.fit(X_scaled)

    results.append({
        "k": k,
        "bic": gmm.bic(X_scaled),
        "aic": gmm.aic(X_scaled)
    })

Plotting BIC

results_df.plot(x='k', y='bic')

The optimal model often appears near the minimum BIC.

Manual Likelihood Verification

One of the best ways to truly understand sklearn is to reconstruct its calculations manually.

Computing Component Densities

from scipy.stats import multivariate_normal

component_densities = np.column_stack([

    multivariate_normal.pdf(
        X_scaled,
        mean=means[j],
        cov=covariances[j],
        allow_singular=True
    )

    for j in range(k)
])

Weighted Densities

\[\pi_k \mathcal{N}(x_i \mid \mu_k,\Sigma_k)\] 
weighted_densities = component_densities * weights

Marginal Likelihood

\[p(x_i)=\sum_k \pi_k \mathcal{N}(x_i|\mu_k,\Sigma_k)\] 
likelihoods = component_densities @ weights

Recovering Responsibilities

probas_manual = (
    weighted_densities
    / likelihoods[:, None]
)

This should match:

gmm.predict_proba(X_scaled)

exactly.

Parameter Calculations, AIC, and BIC

A Gaussian Mixture Model does not just estimate clusters. It also estimates a large number of parameters. That matters because model selection criteria like AIC and BIC penalize models that become too flexible.

For a GMM with:

\[K = \text{number of mixture components}\] \[d = \text{number of features}\] \[n = \text{number of observations}\]

the total number of estimated parameters is:

\[p = (K - 1) + Kd + K\frac{d(d+1)}{2}\]

This comes from three parts.

First, the mixture weights:

\[K - 1\]

There are \(K\) weights, but they must sum to one:

\[\sum_{k=1}^{K}\pi_k = 1\]

So only \(K - 1\) of them are freely estimated.

Second, the means:

\[Kd\]

Each component has one mean vector:

\[\mu_k \in \mathbb{R}^d\]

So across \(K\) components, there are \(Kd\) mean parameters.

Third, the covariance matrices:

\[K\frac{d(d+1)}{2}\]

Each full covariance matrix is symmetric, so instead of estimating \(d^2\) entries, only the diagonal plus the upper or lower triangle must be estimated.

That gives:

\[\frac{d(d+1)}{2}\]

covariance parameters per component.

So the full parameter count is:

\[p = (K - 1) + Kd + K\frac{d(d+1)}{2}\]

In code, this appears as:

p = k - 1 + k*d + k*((d**2 - d)/2 + d)

The covariance part:

((d**2 - d)/2 + d)

is equivalent to:

\[\frac{d(d+1)}{2}\]

because:

\[\frac{d^2-d}{2}+d = \frac{d^2-d+2d}{2} = \frac{d^2+d}{2} = \frac{d(d+1)}{2}\]

So the parameter formula is correct for a GMM with full covariance matrices.

Once \(p\) is known, AIC can be computed as:

\[AIC = 2p - 2\log L\]

and BIC as:

\[BIC = p\log(n) - 2\log L\]

In code:

aic = 2 * p - 2 * log_likelihoods.sum()
bic = p * np.log(n) - 2 * log_likelihoods.sum()

Here:

log_likelihoods.sum()

is the total log likelihood:

\[\log L = \sum_{i=1}^{n}\log p(x_i)\]

where:

\[p(x_i) = \sum_{k=1}^{K} \pi_k \mathcal{N}(x_i \mid \mu_k, \Sigma_k)\]

The likelihood can be computed manually with:

component_densities = np.column_stack([
    multivariate_normal.pdf(
        X,
        mean=means[j],
        cov=covariances[j],
        allow_singular=True
    )
    for j in range(k)
])

weighted_densities = component_densities * weights
likelihoods = component_densities @ weights
log_likelihoods = np.log(likelihoods)

The quantity:

np.mean(log_likelihoods)

matches:

final_gmm.score(X_scaled)

because score returns the average log likelihood per observation.

Meanwhile:

log_likelihoods.sum()

is the quantity used inside AIC and BIC.

Thus, the manual calculations reproduce sklearn’s implementations:

final_gmm.aic(X_scaled)
final_gmm.bic(X_scaled)

AIC and BIC both reward better fit through larger log likelihood, but penalize complexity through \(p\). The difference is that BIC usually penalizes additional parameters more strongly because it uses:

\[p\log(n)\]

instead of:

\[2p\]

So in practice:

  • lower AIC indicates better expected predictive fit

  • lower BIC indicates better fit after a stronger simplicity penalty

  • if AIC and BIC disagree, BIC often selects fewer mixture components

  • comparing AIC and BIC across different values of \(K\) helps justify the selected number of hidden latent profiles or clusters

Visualizing Hidden Profiles with PCA

from sklearn.decomposition import PCA

pca = PCA(n_components=2)

X_pca = pca.fit_transform(X_scaled)

Scatter Plot

plt.scatter(
    X_pca[:,0],
    X_pca[:,1],
    c=df["sleep_profile"]
)

This projects high-dimensional latent structure into 2D.

Interpreting Profiles

We can summarize cluster means:

profile_summary = (
    df.groupby("sleep_profile")[features]
      .mean()
)

This reveals latent physiological states.

For example:

Profile Characteristics
High Recovery low RHR, high HRV
Sleep Debt low sleep duration
Fragmented Sleep many disturbances

Heatmaps for Interpretation

plt.imshow(profile_z, aspect="auto")

Heatmaps make latent-variable interpretation dramatically easier.

Why EM Is Powerful

EM solves problems where:

  • labels are missing

  • latent structure exists

  • probabilistic assignment matters

Unlike K-Means:

K-Means GMM
hard clusters soft clusters
Euclidean distance probabilistic density
spherical assumption full covariance
no uncertainty uncertainty modeling

K-Means as a Special Case

K-Means can actually be derived from GMMs when:

\[\Sigma_k = \sigma^2 I\]

and

\[\sigma^2 \to 0\]

Then soft assignments collapse into hard assignments.

Why Responsibilities Matter

Suppose a datapoint has:

[0.10, 0.85, 0.05]

This means:

  • 10% chance profile 0

  • 85% chance profile 1

  • 5% chance profile 2

That uncertainty is valuable.

Real biology is rarely binary.

Numerical Stability Problems

Real EM implementations face issues:

  • singular covariance matrices

  • underflow

  • exploding likelihoods

  • local optima

Scikit-learn handles these robustly.

That is why reinforcing theory with sklearn is critical.

Complete Libraries Used

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from scipy.stats import multivariate_normal

from sklearn.mixture import GaussianMixture
from sklearn.preprocessing import StandardScaler
from sklearn.impute import SimpleImputer
from sklearn.pipeline import Pipeline
from sklearn.decomposition import PCA

Reference

  1. Yoo, H. J., Nashiro, K., Min, J., Cho, C., Mercer, N., Bachman, S. L., Nasseri, P., Dutt, S., Porat, S., Choi, P., Zhang, Y., Grigoryan, V., Feng, T., Thayer, J. F., Lehrer, P., Chang, C., Stanley, J. A., Head, E., Rouanet, J., … Mather, M. (2021). Heart Rate Variability Biofeedback Training and Emotion Regulation (Version 1.3.5) [Dataset]. OpenNeuro. https://doi.org/10.18112/openneuro.ds003823.v1.3.5