DEV Community: Berkan Sesen

Gaussian Process Regression: The Bayesian Approach to Curve Fitting

Berkan Sesen — Mon, 13 Apr 2026 08:33:12 +0000

You've trained a machine learning model and want to tune its hyperparameters. Each evaluation takes hours. You've tested 6 configurations so far. Where should you try next?

If you read our hyperparameter optimisation post, you saw Bayesian optimisation solve exactly this problem. The secret weapon behind it is a Gaussian process (GP) — a model that predicts not just a value, but how uncertain it is about that value. Near your tested configurations, the GP is confident. Far away, it honestly admits "I don't know."

This is regression with built-in uncertainty quantification. Unlike fitting a line or a polynomial, a GP doesn't commit to a fixed functional form. Instead, it defines a distribution over functions and lets the data narrow it down. The result is a smooth prediction with a confidence band that widens where data is sparse and tightens where it's dense.

By the end of this post, you'll implement GP regression from scratch with NumPy, understand how the kernel function encodes your assumptions about smoothness, and see exactly why GPs power Bayesian optimisation.

Let's Build It

Click the badge to open the interactive notebook:

Watch how the GP posterior sharpens as we feed it observations one by one — the shaded region represents 95% confidence, and it collapses around each data point:

Here's the complete implementation. We'll use 6 noisy observations from Ebden's GP tutorial (2008) and predict across a dense grid:

import numpy as np
from numpy.linalg import inv, det
import matplotlib.pyplot as plt

# --- The Kernel ---
def rbf_kernel(x1, x2, sigma_f, l):
    """Squared exponential (RBF) kernel."""
    return sigma_f**2 * np.exp(-0.5 * (x1 - x2)**2 / l**2)

def kernel_matrix(X1, X2, sigma_f, l):
    """Compute kernel matrix between two sets of points."""
    return np.array([[rbf_kernel(a, b, sigma_f, l) for b in X2] for a in X1])

# --- GP Prediction (Ebden Eq. 8 & 9) ---
def gp_predict(X_train, y_train, X_test, sigma_f, l, sigma_n):
    """Predict mean and variance at test points."""
    n = len(X_train)
    K = kernel_matrix(X_train, X_train, sigma_f, l) + sigma_n**2 * np.eye(n)
    K_s = kernel_matrix(X_test, X_train, sigma_f, l)
    K_ss = kernel_matrix(X_test, X_test, sigma_f, l) + sigma_n**2 * np.eye(len(X_test))

    K_inv = inv(K)
    mu = K_s @ K_inv @ y_train
    cov = K_ss - K_s @ K_inv @ K_s.T

    return mu, cov

# --- Data from Ebden (2008) tutorial ---
X_train = np.array([-1.50, -1.00, -0.75, -0.40, -0.25, 0.00])
y_train = 0.55 * np.array([-3.0, -2.0, -0.6, 0.4, 1.0, 1.6])

# Hyperparameters (optimised values from the tutorial: sigma_f=1.27, l=1.0)
sigma_n = 0.3   # observation noise (known from error bars)
sigma_f = 1.27   # signal standard deviation (optimised)
l = 1.0          # length-scale (optimised)

# Predict on a dense grid
X_test = np.linspace(-2.0, 1.0, 200)
mu, cov = gp_predict(X_train, y_train, X_test, sigma_f, l, sigma_n)
std = np.sqrt(np.diag(cov))

# Plot
fig, ax = plt.subplots(figsize=(10, 5))
ax.fill_between(X_test, mu - 1.96 * std, mu + 1.96 * std,
                alpha=0.2, color='tab:blue', label='95% confidence')
ax.plot(X_test, mu, 'tab:blue', linewidth=2, label='GP mean')
ax.errorbar(X_train, y_train, yerr=sigma_n, fmt='ro', capsize=4,
            markersize=6, label='Training data')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.legend()
plt.show()

The blue line is our best estimate. The shaded band is the 95% confidence interval — wide where we have no data, narrow near the observations. At $x_* = 0.2$ , the GP predicts $\bar{y}_* \approx 0.98$ with variance 0.21, closely matching Ebden's worked example ( $\bar{y}_* = 0.95$ , $\text{var} = 0.21$ ).

What Just Happened?

Three ingredients make GP regression work: a kernel function that encodes our smoothness assumptions, covariance matrices that capture relationships between all points, and conditioning that turns the prior into a posterior.

The Kernel: Encoding Smoothness

The squared exponential (RBF) kernel measures similarity between inputs:

Think of it as answering: "If I know the function value at $x$ , how much does that tell me about $x'$ ?"

Close together ( $|x - x'| \ll \ell$ ): $k \approx \sigma_f^2$ — strong correlation, they "see" each other
Far apart ( $|x - x'| \gg \ell$ ): $k \approx 0$ — independent, no information shared

The two hyperparameters control different things:

$\sigma_f$ (signal std) — the typical amplitude of the function. Higher means larger vertical swings.
$\ell$ (length-scale) — how far you need to move in $x$ before the function value changes significantly. Short $\ell$ gives wiggly functions; long $\ell$ gives smooth ones.

There's also $\sigma_n$ (noise std), capturing measurement noise: each observation $y$ relates to the true function via $y = f(x) + \mathcal{N}(0, \sigma_n^2)$ .

The effect of the length-scale is dramatic:

With $\ell = 0.3$ , the GP tries to wiggle through every point — it over-fits. With $\ell = 3.0$ , it's so smooth it can't follow the data's trend — it under-fits. The optimised $\ell = 1.0$ strikes the right balance.

Building the Covariance Matrices

To make predictions, we compute three kernel matrices:

$K$ — between all training points (Ebden Eq. 4):

The $\sigma_n^2$ on the diagonal accounts for observation noise. Off-diagonal entries capture how correlated two training points are.

$K_*$ — between test and training points (Ebden Eq. 5):

`$K_{}$`** — between test points:

In our code, kernel_matrix computes these. The + sigma_n**2 * np.eye(n) adds noise to the diagonal.

The Prediction Equations

The GP assumption is that training outputs $\mathbf{y}$ and test outputs $y_*$ are jointly Gaussian:

Conditioning on the observed $\mathbf{y}$ gives the posterior — also Gaussian (this is the beauty of Gaussians: conditioning preserves Gaussianity):

The mean $\bar{y}_*$ is a weighted combination of the training outputs, where the weights come from how correlated $x_*$ is with each training point. The variance shrinks where $K_*$ has large entries — near training data — and expands where $x_*$ is far from any observation.

This is Bayesian inference in action: the prior (encoded by the kernel) gets updated by the data to produce a posterior that's both a prediction and an honest assessment of uncertainty.

Going Deeper

GP Prior: What the Kernel Believes Before Seeing Data

Before any observations, the GP prior defines a distribution over functions. We can sample from it by drawing from $\mathcal{N}(\mathbf{0}, K)$ where $K$ is the kernel matrix evaluated on a grid:

X_grid = np.linspace(-2, 2, 200)
K_prior = kernel_matrix(X_grid, X_grid, sigma_f=1.27, l=1.0)
K_prior += 1e-8 * np.eye(len(X_grid))  # numerical stability

# Draw 3 random functions from the prior
np.random.seed(42)
samples = np.random.multivariate_normal(np.zeros(len(X_grid)), K_prior, size=5)

fig, ax = plt.subplots(figsize=(10, 5))
for i in range(5):
    ax.plot(X_grid, samples[i], label=f'Sample {i+1}')
ax.fill_between(X_grid, -1.96 * 1.27, 1.96 * 1.27, alpha=0.1, color='grey',
                label='95% prior band')
ax.set_xlabel('x')
ax.set_ylabel('f(x)')
ax.legend()
plt.show()

Every sample is smooth — the RBF kernel enforces this. The length-scale $\ell = 1.0$ means the functions vary on a scale of roughly 1 unit in $x$ . Data will prune this infinite family down to the functions consistent with our observations.

Hyperparameter Optimisation: Marginal Log-Likelihood

How did we arrive at $\sigma_f = 1.27$ and $\ell = 1.0$ ? The original code optimises the marginal log-likelihood (Ebden Eq. 10):

This balances three terms:

Data fit ( $-\frac{1}{2}\mathbf{y}^\top K^{-1}\mathbf{y}$ ) — how well the model explains the observations
Complexity penalty ( $-\frac{1}{2}\log|K|$ ) — penalises overly flexible models (Occam's razor)
Normalisation ( $-\frac{n}{2}\log 2\pi$ ) — constant

This is maximum likelihood estimation applied to the marginal likelihood — "marginal" because we've integrated out the function values, leaving only the hyperparameters.

from scipy.optimize import minimize

def neg_log_marginal_likelihood(params, X_train, y_train, sigma_n):
    """Negative log marginal likelihood (Ebden Eq. 10)."""
    sigma_f, l = np.abs(params)  # ensure positive
    n = len(X_train)
    K = kernel_matrix(X_train, X_train, sigma_f, l) + sigma_n**2 * np.eye(n)

    log_lik = (-0.5 * y_train @ inv(K) @ y_train
               - 0.5 * np.log(det(K))
               - n / 2 * np.log(2 * np.pi))
    return -log_lik

# Optimise sigma_f and l, keeping sigma_n=0.3 fixed (matching original code)
result = minimize(
    neg_log_marginal_likelihood,
    x0=[0.1, 0.5],  # initial values from original code
    args=(X_train, y_train, 0.3),
    method='Nelder-Mead'
)
sigma_f_opt, l_opt = np.abs(result.x)
print(f'Optimised: sigma_f={sigma_f_opt:.2f}, l={l_opt:.2f}')
# Optimised: sigma_f=1.34, l=1.04

The log-likelihood landscape shows a clear optimum:

The optimiser starts at $(\sigma_f, \ell) = (0.1, 0.5)$ and converges to $(1.34, 1.04)$ , close to the tutorial's reported values of $(1.27, 1.0)$ . The small difference comes from Nelder-Mead finding a slightly different local optimum — the negative log-likelihoods differ by only 0.004. The complexity penalty prevents overfitting — without it, the model would shrink $\ell$ to interpolate every noisy observation exactly.

A fully Bayesian approach would place priors on $\sigma_f$ and $\ell$ and integrate over them using MCMC. For most applications, the point estimate from marginal likelihood optimisation works well.

Validating with sklearn

Let's confirm our from-scratch implementation matches scikit-learn:

from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, WhiteKernel

kernel = 1.27**2 * RBF(length_scale=1.0) + WhiteKernel(noise_level=0.3**2)
gpr = GaussianProcessRegressor(kernel=kernel, optimizer=None, alpha=0)

gpr.fit(X_train.reshape(-1, 1), y_train)
mu_sk, std_sk = gpr.predict(X_test.reshape(-1, 1), return_std=True)

print(f'Max difference in mean: {np.max(np.abs(mu - mu_sk)):.6f}')
# Max difference in mean: ~1e-14 (numerical precision)

Our NumPy implementation produces identical results.

When NOT to Use Gaussian Processes

GPs have a fundamental limitation: cubic scaling. Computing $K^{-1}$ is $O(n^3)$ in the number of training points. With 100 points it's instant; with 1,000 it takes seconds; with 10,000 it becomes impractical.

For large datasets, consider:

Sparse GP approximations — use a subset of inducing points to approximate the full GP
Random Fourier features — approximate the kernel with explicit feature maps
Neural networks — scale linearly in data but lose calibrated uncertainty

The RBF kernel also assumes stationarity — the same smoothness everywhere. If your function is smooth in one region and jagged in another, you'd need a non-stationary kernel or a different model.

Where This Comes From

Rasmussen & Williams (2006)

The definitive reference is Rasmussen, C.E. & Williams, C.K.I. (2006) Gaussian Processes for Machine Learning, MIT Press. The full text is available free online.

They define a Gaussian process formally:

"A Gaussian process is a collection of random variables, any finite number of which have a joint Gaussian distribution."

A GP is fully specified by its mean function $m(x)$ and covariance function $k(x, x')$ :

We set $m(x) = 0$ (standard practice — the kernel handles everything), so the covariance function alone defines the GP. This is the function-space view: we're placing a prior directly on functions, not on parameters.

Rasmussen & Williams also present the weight-space view (Chapter 2.1). In Bayesian linear regression with basis functions $\phi(x)$ , the predictive distribution is:

As the number of basis functions grows to infinity, this converges to the GP formulation — the function-space and weight-space views are equivalent. This connects GPs to neural networks: Neal (1996) showed that a single-hidden-layer neural network with infinitely many hidden units converges to a GP.

Ebden (2008) — Our Implementation Reference

Our implementation follows Ebden, M. (2008) "Gaussian Processes for Regression: A Quick Introduction", which maps directly to our code:

Tutorial Equation	Our Code	What it computes
Eq. (1): RBF kernel	`rbf_kernel()`	$k(x, x') = \sigma_f^2 \exp(-(x-x')^2 / 2\ell^2)$
Eq. (4): $K$ matrix	`kernel_matrix() + sigma_n^2 * I`	Training covariance
Eq. (5): $K_$ and $K_{*}$	`kernel_matrix()` calls	Test-train and test-test covariance
Eq. (8): Predictive mean	`K_s @ K_inv @ y_train`	$\bar{y}_* = K_* K^{-1}\mathbf{y}$
Eq. (9): Predictive variance	`K_ss - K_s @ K_inv @ K_s.T`	$K_{*} - K_ K^{-1} K_*^\top$
Eq. (10): Log-likelihood	`neg_log_marginal_likelihood()`	Hyperparameter optimisation

Historical Context

The idea of using stochastic processes for interpolation has deep roots:

Kolmogorov (1941) and Wiener (1949) — optimal linear prediction theory
Matheron (1963) — "kriging" in geostatistics, named after D.G. Krige's mining work in South Africa
O'Hagan (1978) — formalised the Bayesian interpretation
MacKay (1998) — introduced GPs to the machine learning community
Neal (1996) — proved the GP-neural network connection
Rasmussen & Williams (2006) — the modern comprehensive treatment

Kriging and GP regression are mathematically identical — the geostatistics and machine learning communities developed the same ideas independently, using different vocabulary (variograms vs kernels, kriging variance vs posterior variance).

Try It Yourself

The interactive notebook includes exercises:

Different kernels — Implement the Matern 5/2 kernel and compare its predictions to the RBF. How do the confidence bands differ?
Multi-dimensional inputs — Extend the GP to 2D inputs. The kernel becomes $k(\mathbf{x}, \mathbf{x}') = \sigma_f^2 \exp(-\|\mathbf{x} - \mathbf{x}'\|^2 / 2\ell^2)$ .
Noisy function — Generate data from $y = \sin(x) + \epsilon$ with $\epsilon \sim \mathcal{N}(0, 0.2^2)$ . Fit a GP and observe how the confidence band captures the noise.
Bayesian optimisation — Use your GP to implement a simple acquisition function (Expected Improvement) and optimise a 1D black-box function.

Interactive Tools

GP Regression Playground — Fit Gaussian processes to your own data and experiment with kernels in the browser

Hyperparameter Optimization: Grid vs Random vs Bayesian — See GPs as the surrogate model in Bayesian optimisation
From MLE to Bayesian Inference — The Bayesian framework that underpins GP regression
MCMC Island Hopping: Understanding Metropolis-Hastings — An alternative to point estimates for GP hyperparameters

Frequently Asked Questions

What does a Gaussian process actually model?

A Gaussian process defines a probability distribution over functions, not just over individual predictions. Any finite collection of function values is modelled as a multivariate Gaussian distribution. The kernel function specifies how correlated any two function values are, which determines the smoothness and structure of the functions the GP considers plausible.

What does the length-scale hyperparameter control?

The length-scale determines how quickly the function can vary. A short length-scale allows rapid, wiggly changes and can lead to overfitting, while a long length-scale enforces slow, smooth variation and can underfit. The optimal length-scale is typically found by maximising the marginal log-likelihood, which automatically balances data fit against model complexity.

Why do Gaussian processes scale poorly to large datasets?

GP prediction requires inverting the training covariance matrix, which is an O(n^3) operation. For 10,000 or more training points, this becomes computationally prohibitive. Sparse GP approximations, which use a smaller set of inducing points to approximate the full covariance, are the most common workaround.

How is GP regression related to Bayesian optimisation?

In Bayesian optimisation, a GP serves as a surrogate model that approximates the expensive objective function. The GP's uncertainty estimates are critical: an acquisition function uses the predicted mean and variance to decide where to evaluate next, balancing exploitation (areas with good predicted values) and exploration (areas with high uncertainty).

Can I use kernels other than the RBF?

Yes. The choice of kernel encodes your assumptions about the function. The Matern kernel allows you to control the differentiability of the function (the RBF assumes infinite differentiability). Periodic kernels capture repeating patterns. You can also combine kernels by adding or multiplying them to model functions with multiple structural components.

Hyperparameter Optimization: Grid vs Random vs Bayesian

Berkan Sesen — Fri, 10 Apr 2026 08:20:40 +0000

You've trained a Random Forest and it works — 85% accuracy out of the box. But you used the default hyperparameters. What if n_estimators=500 with max_features=0.3 and min_samples_leaf=10 pushes that to 91%? Only one way to find out: search.

The problem is combinatorial. Our Random Forest has 4 hyperparameters. If you try 10 values for each in a grid, that's $10^4 = 10{,}000$ combinations. Each combination requires 5-fold cross-validation. That's 50,000 model fits — and we only have 4 dimensions. Neural networks routinely have 10–20 tunable hyperparameters, where exhaustive search is physically impossible.

This post compares three strategies of increasing sophistication:

Grid Search — try every combination on a predefined grid
Random Search — sample combinations at random (surprisingly effective)
Bayesian Optimization — build a model of the objective and use it to choose the next point intelligently

We'll run all three on the same classification task, using the same Random Forest and the same hyperparameter ranges. You'll see that for an easy problem, all three reach ~90% accuracy — but the way they get there reveals fundamentally different philosophies about search. The real payoff of Bayesian optimization comes when evaluations are expensive: training a neural network for hours, running a simulation, or calling a paid API.

If you've read the genetic algorithms post, you've already seen one approach to gradient-free optimization. Hyperparameter optimization is another — and Bayesian optimization is arguably the most elegant solution.

Quick Win: Run All Three Methods

Click the badge to open the notebook and run everything yourself:

Setup: The Problem

We'll classify a synthetic dataset with 2,000 samples, 20 features, and 4 classes — a moderately challenging multi-class problem:

import numpy as np
from sklearn.datasets import make_classification
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score, GridSearchCV, RandomizedSearchCV, StratifiedKFold
from sklearn import metrics
from skopt import gp_minimize
from skopt.space import Real, Integer, Categorical

# Generate dataset (tuned to match the original Kaggle mobile price dataset's ~90% RF accuracy)
X, y = make_classification(
    n_samples=2000, n_features=20, n_informative=15,
    n_clusters_per_class=1, flip_y=0.01,
    n_classes=4, random_state=42
)

print(f'{X.shape[1]} features, {len(np.unique(y))} classes, {X.shape[0]} samples')
# 20 features, 4 classes, 2000 samples

All three methods will tune the same 4 hyperparameters with 5-fold stratified cross-validation:

Hyperparameter	Type	Range	What it controls
`max_features`	Continuous	[0.1, 1.0]	Fraction of features per split
`n_estimators`	Integer	[100, 1000]	Number of trees
`min_samples_leaf`	Integer	[5, 25]	Minimum samples in a leaf
`criterion`	Categorical	{gini, entropy}	Split quality measure

Method 1: Grid Search

Grid search evaluates every combination on a predefined grid:

param_grid = {
    'n_estimators': [200, 400, 600, 800],
    'criterion': ['gini', 'entropy'],
    'min_samples_leaf': [5, 10, 20],
    'max_features': [0.3, 0.5, 0.8]
}

grid_search = GridSearchCV(
    estimator=RandomForestClassifier(n_jobs=-1, random_state=42),
    param_grid=param_grid,
    scoring='accuracy',
    cv=StratifiedKFold(n_splits=5, shuffle=True, random_state=42),
    verbose=0,
    n_jobs=-1
)

grid_search.fit(X, y)
print(f'Grid Search — Best accuracy: {grid_search.best_score_:.4f}')
print(f'Combinations evaluated: {len(grid_search.cv_results_["mean_test_score"])}')
print(f'Best params: {grid_search.best_params_}')

With $4 \times 2 \times 3 \times 3 = 72$ combinations and 5 folds each, that's 360 model fits.

Method 2: Random Search

Random search samples 15 combinations uniformly at random from the ranges:

param_distributions = {
    'n_estimators': np.arange(100, 1001),
    'criterion': ['gini', 'entropy'],
    'min_samples_leaf': np.arange(5, 26),
    'max_features': np.linspace(0.1, 1.0, 100)
}

random_search = RandomizedSearchCV(
    estimator=RandomForestClassifier(n_jobs=-1, random_state=42),
    param_distributions=param_distributions,
    n_iter=15,
    scoring='accuracy',
    cv=StratifiedKFold(n_splits=5, shuffle=True, random_state=42),
    verbose=0,
    n_jobs=-1,
    random_state=42
)

random_search.fit(X, y)
print(f'Random Search — Best accuracy: {random_search.best_score_:.4f}')
print(f'Combinations evaluated: 15')
print(f'Best params: {random_search.best_params_}')

Only 15 combinations, 75 model fits — a fraction of grid search.

Method 3: Bayesian Optimization (Gaussian Process)

Bayesian optimization builds a probabilistic model of the objective and uses it to decide where to evaluate next:

def evaluate_params(params):
    """5-fold stratified CV accuracy for a RandomForest configuration.
    Returns negative accuracy (since gp_minimize minimises)."""
    max_features, n_estimators, min_samples_leaf, criterion = params
    model = RandomForestClassifier(
        max_features=max_features,
        n_estimators=n_estimators,
        min_samples_leaf=min_samples_leaf,
        criterion=criterion,
        n_jobs=-1,
        random_state=42
    )
    kf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    accuracies = []
    for train_idx, val_idx in kf.split(X, y):
        model.fit(X[train_idx], y[train_idx])
        preds = model.predict(X[val_idx])
        accuracies.append(metrics.accuracy_score(y[val_idx], preds))
    return -np.mean(accuracies)

param_space = [
    Real(0.1, 1.0, prior='uniform', name='max_features'),
    Integer(100, 1000, name='n_estimators'),
    Integer(5, 25, name='min_samples_leaf'),
    Categorical(['gini', 'entropy'], name='criterion')
]

result = gp_minimize(
    evaluate_params,
    dimensions=param_space,
    n_calls=15,
    n_random_starts=10,
    random_state=42,
    verbose=False
)

best_params = dict(zip(
    ['max_features', 'n_estimators', 'min_samples_leaf', 'criterion'],
    result.x
))
print(f'Bayesian (GP) — Best accuracy: {-result.fun:.4f}')
print(f'Evaluations: 15 (10 random + 5 guided)')
print(f'Best params: {best_params}')

Same 15 evaluations, but the last 5 are informed by the GP model built from the first 10.

Results Comparison

Method	Best CV Accuracy	Evaluations	Strategy
Grid Search	~90%	72	Exhaustive (predefined grid)
Random Search	~90%	15	Uniform random sampling
Bayesian (GP)	~90%	15 (10 random + 5 guided)	Model-based sequential

All three methods converge to approximately the same accuracy — around 90%. For this well-behaved 4-dimensional problem, the objective landscape is smooth enough that random search finds good regions quickly.

The convergence chart tells the real story. Grid search plods through its predefined grid systematically. Random search jumps around but finds a good region early. Bayesian optimization starts random (first 10 points) then accelerates — the GP model narrows in on promising regions.

The punchline: the methods differ most when evaluations are expensive. For a Random Forest on 2,000 samples, each evaluation takes under a second and you can afford to be wasteful. Train a Transformer for 8 hours per evaluation, and the difference between 15 smart evaluations and 72 brute-force ones is the difference between two days and two weeks.

What Just Happened?

All three methods solve the same problem — find the hyperparameter combination that maximises cross-validated accuracy — but they navigate the search space in fundamentally different ways.

Grid Search: The Cartesian Product

Grid search evaluates every point on a regular lattice. If you specify 4 values for n_estimators, 2 for criterion, 3 for min_samples_leaf, and 3 for max_features, you get $4 \times 2 \times 3 \times 3 = 72$ combinations. It's the Cartesian product of your parameter lists.

Think of a tourist exploring a city by visiting every intersection on the street grid. Thorough? Yes. Efficient? Not even close. The problem is that most intersections are uninteresting — the accuracy surface for a Random Forest is usually smooth, so neighbouring grid points give nearly identical results.

The deeper issue is the curse of dimensionality. In $d$ dimensions with $k$ values per dimension, the grid has $k^d$ points. With 10 values per dimension: 2D gives 100 points (fine), 4D gives 10,000 (slow), 10D gives 10 billion (impossible). Grid search is fundamentally limited to low-dimensional problems with coarse grids.

Random Search: The Bergstra-Bengio Insight

Random search replaces the grid with uniform random samples. This seems like it should be worse — throwing darts blindfolded. But Bergstra and Bengio (2012) showed it's almost always better than grid search for the same computational budget.

The insight is elegant. Suppose only 2 of your 4 hyperparameters actually matter (a common situation). On a $4 \times 4$ grid in 2D, you get 16 unique points — but only 4 unique values per dimension. Random search with 16 points gives you 16 unique values per dimension. You're covering the important dimensions far more densely.

The figure makes this vivid. Grid search evaluates 9 points, but only 3 unique values along each axis. Random search with 9 points explores 9 unique values per axis. Bayesian optimization clusters its samples in the high-accuracy region (top-right), exploring broadly first and then exploiting the best area.

Bayesian Optimization: Learning from History

Grid and random search are memoryless — each evaluation is independent of the others. The 15th evaluation in random search knows nothing about the first 14.

Bayesian optimization is fundamentally different. It maintains a model of the objective function — a Gaussian process (GP) that predicts what the accuracy will be at any point in the hyperparameter space, along with an uncertainty estimate. After each evaluation, the model updates, and an acquisition function decides where to evaluate next.

The three components:

Surrogate model (Gaussian Process) — A probabilistic model that interpolates between observed points. At observed points, uncertainty is zero. Far from observations, uncertainty is high. Think of it as drawing a smooth surface through scattered data points, with error bars that widen between points.
Acquisition function — A cheap-to-evaluate function that balances exploration (high uncertainty) and exploitation (high predicted value). The most common is Expected Improvement (EI): "how much better than the current best do I expect this point to be?"
Optimize the acquisition function — Find the point that maximises EI (a standard optimization problem, but cheap because the GP is fast to evaluate), then run the expensive evaluation there.

This loop — fit GP, maximize acquisition, evaluate, repeat — is what makes Bayesian optimization sequential and adaptive. If you've read the Bayesian inference post, you'll recognise the pattern: start with a prior (the GP), observe data, update to a posterior, and use the posterior to make decisions. The GP generalises this from updating parameter estimates to updating an entire function estimate.

Going Deeper

The GP Surrogate in Action

The animation shows Bayesian optimization on a 1D function. The blue line is the true (unknown) objective. The black line is the GP's mean prediction, with the shaded region showing the 95% confidence interval. The green area below is the Expected Improvement — it peaks where the model expects to find values better than the current best.

Watch how the GP evolves:

Early frames: Few observations, wide uncertainty, EI spreads across unexplored regions (exploration)
Middle frames: The GP starts to capture the function's shape, uncertainty shrinks near data points
Late frames: EI concentrates around the global optimum as the model becomes confident elsewhere (exploitation)

GP Predictive Distribution

At any unobserved point $\mathbf{x}_*$ , the GP predicts a Gaussian distribution:

Where:

$\mu(\mathbf{x}_*)$ — the mean prediction, interpolating nearby observations
$\sigma^2(\mathbf{x}_*)$ — the predictive variance: high far from data, zero at observed points
The predictions are conditioned on all previous observations $\{(\mathbf{x}_i, y_i)\}_{i=1}^n$

The mean and variance are computed in closed form from the kernel function (typically Matérn 5/2 in skopt):

Where $\mathbf{K}$ is the kernel matrix between all observed points, $\mathbf{k}_*$ is the kernel vector between $\mathbf{x}_*$ and observed points, and $\sigma_n^2$ is the observation noise.

In plain English: the GP prediction at a new point is a weighted average of observed values, where the weights come from how "similar" (in kernel space) the new point is to each observation.

Acquisition Functions: Deciding Where to Look Next

The acquisition function turns the GP's prediction into a decision: where should we evaluate next? Three common choices:

Expected Improvement (EI) — the default in skopt:

Where $f(\mathbf{x}^+)$ is the best value observed so far. EI asks: "in expectation, how much will this point improve on the current best?" Points with high predicted mean and high uncertainty score well — this naturally balances exploitation and exploration.

Lower Confidence Bound (LCB):

Minimize the mean minus a multiple of the standard deviation. The parameter $\kappa$ controls the exploration–exploitation trade-off directly: higher $\kappa$ means more exploration.

Probability of Improvement (PI):

The probability that a point improves on the current best. Simpler than EI, but tends to exploit too aggressively — it doesn't care how much better a point might be, only whether it's better at all.

Exploration vs Exploitation

This tension appears everywhere in machine learning. In Q-learning, epsilon-greedy balances trying new actions (exploration) with choosing the best-known action (exploitation). In MCMC, the proposal distribution must explore the parameter space while spending time in high-probability regions.

Bayesian optimization handles this automatically through the acquisition function. EI naturally favours:

Points with high predicted accuracy (exploitation)
Points with high uncertainty (exploration)

No manual schedule needed — the GP's uncertainty estimates do the work.

When NOT to Use Bayesian Optimization

Bayesian optimization isn't always the right tool:

Cheap evaluations — If each evaluation takes seconds (like our Random Forest), random search with 100 iterations is simpler and nearly as effective
High dimensions ( $d > 20$ ) — GPs scale poorly with dimensionality. The kernel becomes uninformative and the acquisition function has too many local optima
Massive parallelism — If you have 1,000 GPUs, you can evaluate 1,000 random configurations simultaneously. Bayesian optimization is inherently sequential (each evaluation depends on all previous ones)
Discrete/conditional spaces — GPs assume smooth, continuous objectives. Deeply nested conditional hyperparameters (e.g., "layer type" → "layer-specific params") are better handled by tree-based methods like Optuna's TPE

Rule of thumb: Use Bayesian optimization when each evaluation costs minutes to hours (neural network training, simulation runs, expensive API calls) and you're in 5–15 dimensions.

Hyperparameters

All values come directly from the original code in quant_code/python/HyperParameter_Optimization/src/:

Parameter	Source file	Value	Why
`max_features`	`gp_min_optim.py`	Real(0.1, 1.0)	Fraction of features per split
`n_estimators`	`gp_min_optim.py`	Integer(100, 1000)	Number of trees in the forest
`min_samples_leaf`	`gp_min_optim.py`	Integer(5, 25)	Minimum leaf samples (regularisation)
`criterion`	`grid_n_random_search.py`	{gini, entropy}	Split quality metric
`n_calls`	`gp_min_optim.py`	15	Total evaluations for GP
`n_random_starts`	`gp_min_optim.py`	10	Initial random exploration
`cv`	all files	5-fold stratified	Evaluation protocol

Deep Dive: The Papers

Bergstra & Bengio (2012) — Random Search for Hyper-Parameter Optimization

The key paper that changed how practitioners think about hyperparameter search. Random Search for Hyper-Parameter Optimization (JMLR) demonstrated that random search is more efficient than grid search for most practical problems.

The core theorem is deceptively simple. Define the effective dimensionality of a search problem as the number of hyperparameters that significantly affect performance. Bergstra and Bengio showed:

"Random search is more efficient than grid search because it allows each trial to explore a different value of every hyperparameter. For problems with low effective dimensionality, this results in dramatically better coverage of the important dimensions."

Their famous figure (reproduced in our sampling comparison above) shows a 2D search where only one dimension matters. Grid search with 9 points wastes 6 of them evaluating the same 3 values of the important dimension. Random search with 9 points gives 9 unique values along the important dimension.

The practical implication: for the same computational budget, random search achieves equal or better results than grid search in virtually all cases. There is essentially no scenario where grid search is preferable.

Snoek, Larochelle & Adams (2012) — Practical Bayesian Optimization

Practical Bayesian Optimization of Machine Learning Algorithms (NeurIPS) introduced the Spearmint system and demonstrated that Bayesian optimization could match or beat human experts at tuning neural networks.

Their algorithm:

Input: Search space S, objective function f, budget N
1. Initialize with n₀ random evaluations
2. For i = n₀+1 to N:
   a. Fit GP to all observations {(x₁,y₁), ..., (xᵢ₋₁,yᵢ₋₁)}
   b. Find xᵢ = argmax_x EI(x) using the GP
   c. Evaluate yᵢ = f(xᵢ)
3. Return x with best observed y

This is exactly what skopt.gp_minimize implements. The key contributions beyond earlier work:

Automatic Relevance Determination (ARD) kernels — learns which hyperparameters matter most
Integrated acquisition function — marginalises over GP hyperparameters rather than fixing them
Pending evaluations — supports parallel evaluations through "fantasised" outcomes

Snoek et al. demonstrated that Bayesian optimization with 30 evaluations outperformed a human expert who had months to tune the same neural networks. The quote that launched a thousand AutoML papers:

"Bayesian optimization spent 2 minutes of overhead in exchange for saving the researchers days of manual tuning."

Mockus et al. (1978) & Jones et al. (1998) — The Origins

Bayesian optimization predates machine learning. Mockus, Tiesis, and Žilinskas (1978) formalised the idea of using a Bayesian model to guide sequential optimization in their work on Bayesian Methods for Seeking the Extremum. They introduced the Expected Improvement criterion, proving it is the optimal policy for a one-step lookahead under certain assumptions.

Two decades later, Jones, Schonlau, and Welch (1998) brought these ideas to engineering design optimization with the EGO (Efficient Global Optimization) algorithm. Their paper Efficient Global Optimization of Expensive Black-Box Functions established the GP + EI framework that Snoek et al. later applied to ML hyperparameters.

Historical Timeline

Mockus et al. (1978) — Bayesian optimization and Expected Improvement
Jones et al. (1998) — EGO: GP + EI for engineering design
Hutter et al. (2011) — SMAC: random forests as surrogate instead of GPs (scales to high dimensions)
Bergstra et al. (2011) — Hyperopt and TPE: tree-structured Parzen estimators (handles conditional spaces)
Bergstra & Bengio (2012) — The random search paper
Snoek et al. (2012) — Practical Bayesian optimization (Spearmint)
Akiba et al. (2019) — Optuna: define-by-run API, pruning, modern TPE (now the most popular framework)

Modern Alternatives: Optuna and Hyperopt

The original source code also includes implementations in Optuna and Hyperopt. Both use Tree-structured Parzen Estimators (TPE) instead of Gaussian Processes:

# Optuna (from optuna_optim.py) — define-by-run API
import optuna

def objective(trial):
    criterion = trial.suggest_categorical('criterion', ['gini', 'entropy'])
    n_estimators = trial.suggest_int('n_estimators', 100, 1000)
    min_samples_leaf = trial.suggest_int('min_samples_leaf', 5, 25)
    max_features = trial.suggest_float('max_features', 0.1, 1.0)

    model = RandomForestClassifier(
        n_estimators=n_estimators, criterion=criterion,
        max_features=max_features, min_samples_leaf=min_samples_leaf,
        n_jobs=-1, random_state=42
    )
    kf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    scores = []
    for train_idx, val_idx in kf.split(X, y):
        model.fit(X[train_idx], y[train_idx])
        scores.append(metrics.accuracy_score(y[val_idx], model.predict(X[val_idx])))
    return -np.mean(scores)

study = optuna.create_study(direction='minimize')
study.optimize(objective, n_trials=15)

TPE models $P(\mathbf{x} | y < y^*)$ and $P(\mathbf{x} | y \geq y^*)$ separately (two Parzen estimators), then selects points that maximise the ratio. This handles conditional and discrete parameters more naturally than GPs.

When to use which:

skopt (GP) — smooth, low-dimensional spaces; small evaluation budgets
Optuna (TPE) — large search spaces; conditional parameters; early pruning of bad trials
Hyperopt (TPE) — similar to Optuna, but older API; still widely used

Interactive Tools

Regression Playground — Experiment with model complexity and see how different hyperparameters affect the fit
Overfitting Explorer — Visualise the bias-variance tradeoff that hyperparameter tuning aims to optimise

Genetic Algorithms from Scratch — Another gradient-free optimizer for black-box functions, using evolutionary strategies instead of surrogate models
From MLE to Bayesian Inference — The Gaussian Process generalises Bayesian updating from parameters to entire functions
MCMC Island Hopping — Exploration vs exploitation in sampling — the same tension that drives acquisition functions
MLE from Scratch — Cross-validated accuracy is a likelihood proxy: we're maximising the probability that the model explains held-out data

Try It Yourself

The interactive notebook includes exercises:

More iterations — Increase n_calls to 50 for the Bayesian optimizer. Does the extra budget find meaningfully better configurations, or does it plateau?
Higher dimensions — Add max_depth (Integer, 5–50) and min_samples_split (Integer, 2–20) as hyperparameters. How do the methods scale with 6 dimensions?
Optuna comparison — Run the Optuna TPE sampler alongside GP-based optimization. Compare convergence curves. Does TPE find good configurations faster?
Acquisition function sweep — Try acq_func='LCB' and acq_func='PI' in gp_minimize. How does the exploration–exploitation balance change?
Simulated expensive evaluations — Add time.sleep(2) inside evaluate_params to simulate expensive evaluations. Now the wall-clock difference between 15 evaluations (Bayesian) and 72 (grid) becomes tangible

The three methods represent a progression in sophistication: grid search makes no assumptions, random search exploits the low effective dimensionality of most problems, and Bayesian optimization builds a model to make each evaluation count. For cheap models like Random Forests, random search is usually sufficient. But when each evaluation costs real time — training a Transformer, running a physics simulation, querying a paid API — the model-based approach of Bayesian optimization can save days of compute. The cost of fitting a GP is negligible compared to hours of training, and even 5 guided evaluations out of 15 total can find configurations that random search might need 50 iterations to reach.

Frequently Asked Questions

What is hyperparameter optimisation and why does it matter?

Hyperparameters are settings you choose before training a model (learning rate, tree depth, regularisation strength) that cannot be learned from the data. Poor choices lead to underfitting or overfitting. Systematic optimisation finds the best combination, often improving model performance significantly compared to manual tuning.

When should I use random search instead of grid search?

Almost always. Random search is more efficient because it explores more unique values of each hyperparameter. Grid search wastes evaluations on unimportant parameter combinations, especially when only one or two hyperparameters actually matter. Random search achieves the same or better results with fewer evaluations in most practical settings.

What is Bayesian optimisation and when is it worth the overhead?

Bayesian optimisation builds a probabilistic model (typically a Gaussian process) of the objective function and uses it to choose the next hyperparameters to evaluate intelligently. It is worth the overhead when each evaluation is expensive (training takes hours) and the search space has fewer than about 20 dimensions. For cheap evaluations, random search is often sufficient.

How many hyperparameter evaluations should I run?

For random search, 60 evaluations is usually enough to find a good configuration if only 1-3 hyperparameters matter (there is a 95% chance of sampling within the top 5% of the space). For Bayesian optimisation, 20-50 evaluations often suffice due to the intelligent search strategy. Scale up for higher-dimensional spaces.

Should I use cross-validation during hyperparameter optimisation?

Yes. Evaluating hyperparameters on a single train-test split is noisy and can lead to overfitting the validation set. K-fold cross-validation gives more reliable performance estimates. Use 5-fold as a default, or 3-fold if training is expensive. Always keep a final held-out test set that is never used during optimisation.

Policy Gradients: REINFORCE from Scratch with NumPy

Berkan Sesen — Wed, 08 Apr 2026 07:54:50 +0000

In the DQN post, we trained a neural network to estimate Q-values and then picked the best action with argmax. That works when the action space is discrete — push left or push right. But what if you need to control a robotic arm with continuous joint angles, or steer a car with a continuous throttle? You can't argmax over infinity.

Policy gradient methods flip the approach: instead of learning a value function and deriving a policy, we directly parameterise the policy and optimise it via gradient ascent. The network outputs action probabilities, we sample from them, and we nudge the parameters toward actions that led to high rewards. No Q-values, no argmax, no experience replay — just a policy, a gradient, and a reward signal.

By the end of this post, you'll implement the REINFORCE algorithm entirely from scratch in NumPy — including the forward pass, backpropagation, and RMSProp optimiser — and train it to balance CartPole. The entire implementation is about 100 lines. No PyTorch, no TensorFlow, just NumPy and the chain rule.

Quick Win: Run the Algorithm

Let's see REINFORCE in action. Click the badge to open the interactive notebook:

The animation shows the agent converging to the maximum score of 500 within about 3,000 episodes, using nothing but NumPy. No deep learning framework, no replay buffer — just a policy, a gradient, and RMSProp.

import numpy as np
import gymnasium as gym

# --- Hyperparameters ---
# Original code (CartPole-v0, max 200 steps): H=100, lr=1e-4, gamma=0.95, batch_size=5
# Adapted for CartPole-v1 (max 500 steps): higher gamma and learning rate
H = 100              # hidden layer neurons
batch_size = 5       # episodes per parameter update
learning_rate = 1e-3 # RMSProp learning rate (original: 1e-4, raised for longer episodes)
gamma = 0.99         # discount factor (original: 0.95, raised for longer horizon)
decay_rate = 0.99    # RMSProp decay
epsilon = 1e-5       # RMSProp epsilon

# --- Network functions ---
def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))

def forward(x, model):
    h = np.dot(model['W1'], x)
    h[h < 0] = 0  # ReLU
    logp = np.dot(model['W2'], h)
    p = sigmoid(logp)
    return p, h

def backward(eph, epdlogp, epx, model):
    dW2 = np.dot(eph.T, epdlogp).ravel()
    dh = np.outer(epdlogp, model['W2'])
    dh[eph <= 0] = 0  # backprop through ReLU
    dW1 = np.dot(dh.T, epx)
    return {'W1': dW1, 'W2': dW2}

def discount_rewards(r, gamma):
    """Standard full-horizon discounting."""
    discounted_r = np.zeros_like(r)
    running_add = 0
    for t in reversed(range(r.size)):
        running_add = running_add * gamma + r[t]
        discounted_r[t] = running_add
    return discounted_r

# --- Initialisation ---
env = gym.make('CartPole-v1')
observation, _ = env.reset()
D = len(observation)  # input dimension (4 for CartPole)

model = {
    'W1': np.random.randn(H, D) / np.sqrt(D),  # Xavier init
    'W2': np.random.randn(H) / np.sqrt(H),
}
grad_buffer = {k: np.zeros_like(v) for k, v in model.items()}
rmsprop_cache = {k: np.zeros_like(v) for k, v in model.items()}

# --- Training loop ---
xs, hs, dlogps, drs, episode_durations = [], [], [], [], []
episode_number = 0
t = 0

while episode_number < 5000:
    x = observation
    aprob, h = forward(x, model)
    action = 1 if np.random.uniform() < aprob else 0

    xs.append(x)
    hs.append(h)
    dlogps.append(action - aprob)  # policy gradient

    observation, reward, terminated, truncated, _ = env.step(action)
    drs.append(reward)

    if terminated or truncated:
        episode_number += 1
        episode_durations.append(t)
        t = 0

        epx = np.vstack(xs)
        eph = np.vstack(hs)
        epdlogp = np.vstack(dlogps)
        epr = np.vstack(drs)
        xs, hs, dlogps, drs = [], [], [], []

        # Discount and standardise rewards
        discounted_epr = discount_rewards(epr, gamma)
        discounted_epr -= np.mean(discounted_epr)
        std = np.std(discounted_epr)
        if std > 0:
            discounted_epr /= std

        # The PG magic: weight gradients by advantage
        epdlogp *= discounted_epr
        grad = backward(eph, epdlogp, epx, model)
        for k in model:
            grad_buffer[k] += grad[k]

        # RMSProp update every batch_size episodes
        if episode_number % batch_size == 0:
            for k, v in model.items():
                g = grad_buffer[k]
                rmsprop_cache[k] = decay_rate * rmsprop_cache[k] + (1 - decay_rate) * g**2
                model[k] += learning_rate * g / (np.sqrt(rmsprop_cache[k]) + epsilon)
                grad_buffer[k] = np.zeros_like(v)

        if episode_number % 500 == 0:
            avg = np.mean(episode_durations[-100:])
            print(f'Episode {episode_number}, 100-ep avg: {avg:.1f}')

        observation, _ = env.reset()
    t += 1

env.close()
print(f'Final 100-episode average: {np.mean(episode_durations[-100:]):.1f}')

The result: The agent converges to the 500-step maximum, using nothing but NumPy. No deep learning framework — we compute every gradient by hand.

Visualise the Learning

import matplotlib.pyplot as plt

rolling = np.convolve(episode_durations, np.ones(100)/100, mode='valid')

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(rolling, 'b-', linewidth=0.8)
ax.axhline(y=500, color='g', linestyle='--', alpha=0.5, label='Max score (500)')
ax.set_xlabel('Episode')
ax.set_ylabel('Duration (100-episode rolling avg)')
ax.set_title('REINFORCE on CartPole-v1')
ax.set_ylim(0, 550)
ax.legend()
ax.grid(True, alpha=0.3)
plt.show()

What Just Happened?

We built a complete RL agent with no frameworks — just a policy network, a gradient, and a reward signal. Let's walk through each piece.

The Policy Network (4 → 100 → 1)

The network is a two-layer perceptron that maps a 4-dimensional CartPole state to a single probability:

State [x, ẋ, θ, θ̇]  →  Hidden (100 ReLU)  →  Output (sigmoid)  →  P(push right)
         4                    100                     1

def forward(x, model):
    h = np.dot(model['W1'], x)   # (100, 4) × (4,) → (100,)
    h[h < 0] = 0                  # ReLU activation
    logp = np.dot(model['W2'], h) # (100,) × (100,) → scalar
    p = sigmoid(logp)             # squash to [0, 1]
    return p, h

The output p is the probability of pushing right. We sample from this Bernoulli distribution:

action = 1 if np.random.uniform() < aprob else 0

This is fundamentally different from DQN's approach. DQN outputs Q-values for every action and picks the highest (deterministic, epsilon-greedy). Here, the network outputs a probability and we sample — the policy is stochastic by design. This built-in exploration means we don't need an epsilon schedule.

The Policy Gradient Signal

After taking an action, we record the "gradient that encourages the action that was taken":

dlogps.append(action - aprob)

If we pushed right (action=1) and aprob=0.7, then dlogp = 1 - 0.7 = 0.3 — a positive gradient that nudges the network to make "push right" even more likely. If we pushed left (action=0) and aprob=0.7, then dlogp = 0 - 0.7 = -0.7 — a negative gradient that decreases the probability of pushing right (making left more likely).

But at this point, we don't know if the action was good. That's what the reward tells us.

The PG Magic Line

This is where policy gradients really happen:

epdlogp *= discounted_epr  # modulate gradient with advantage

Every action's gradient gets multiplied by its discounted reward. Good actions (high reward) get their gradients amplified — "do more of this". Bad actions (low or negative reward) get their gradients flipped — "do less of that". This single line is the heart of REINFORCE.

Think of it like a coach watching game film: every play gets a grade. Plays that led to scoring get reinforced; plays that led to turnovers get discouraged. The magnitude of the grade determines how strongly the feedback applies.

Reward Discounting

CartPole gives +1 reward at every timestep the pole stays up. We discount future rewards so earlier actions, which contributed to a longer run, receive more credit:

def discount_rewards(r, gamma):
    discounted_r = np.zeros_like(r)
    running_add = 0
    for t in reversed(range(r.size)):
        running_add = running_add * gamma + r[t]
        discounted_r[t] = running_add
    return discounted_r

With gamma=0.99, an action taken 100 steps before the end still receives $0.99^{100} \approx 0.37$ of the terminal reward. This gives the network a smooth gradient across the episode: early actions in long episodes receive higher discounted returns than the same actions in short episodes, so the network learns to favour strategies that keep the pole up for longer.

Reward Standardisation (Variance Reduction)

After discounting, we standardise the rewards to have zero mean and unit variance:

discounted_epr -= np.mean(discounted_epr)
std = np.std(discounted_epr)
if std > 0:
    discounted_epr /= std

This is critical for stable training. Without standardisation, the magnitude of the gradient depends on the absolute reward scale. With it, roughly half the actions get reinforced (above-average) and half get discouraged (below-average). It's a simple form of a baseline — one of the most important variance reduction techniques in policy gradients.

Manual Backpropagation

We compute gradients by hand, just as in the backpropagation post:

def backward(eph, epdlogp, epx, model):
    dW2 = np.dot(eph.T, epdlogp).ravel()     # gradient for output weights
    dh = np.outer(epdlogp, model['W2'])        # backprop to hidden layer
    dh[eph <= 0] = 0                           # backprop through ReLU
    dW1 = np.dot(dh.T, epx)                   # gradient for input weights
    return {'W1': dW1, 'W2': dW2}

The chain rule flows backwards: output layer gradient → hidden layer gradient (masked by ReLU) → input layer gradient. This is identical to standard backprop, except the "loss" is the policy gradient signal (epdlogp already weighted by advantage).

Batch Gradient Accumulation

Rather than updating after every episode, we accumulate gradients over batch_size=5 episodes:

for k in model:
    grad_buffer[k] += grad[k]

if episode_number % batch_size == 0:
    # RMSProp update using accumulated gradients
    ...

This reduces gradient variance — each update reflects 5 episodes worth of experience rather than just one. It's the same idea as minibatch gradient descent in supervised learning.

RMSProp from Scratch

The optimiser is RMSProp, implemented manually:

rmsprop_cache[k] = decay_rate * rmsprop_cache[k] + (1 - decay_rate) * g**2
model[k] += learning_rate * g / (np.sqrt(rmsprop_cache[k]) + epsilon)

RMSProp maintains a running average of squared gradients and divides by their root. This gives adaptive per-parameter learning rates: parameters with consistently large gradients get smaller effective learning rates, and vice versa. The decay_rate=0.99 means the running average looks back roughly 100 updates.

Going Deeper

The Policy Gradient Theorem

The goal of policy gradients is to maximise expected cumulative reward:

Where $\tau$ is a trajectory (sequence of states and actions) sampled under policy $\pi_\theta$ . The policy gradient theorem (Sutton et al., 1999) tells us how to compute the gradient of this objective:

Where $G_t = \sum_{k=t}^{T} \gamma^{k-t} r_k$ is the return from timestep $t$ .

The intuition: $\nabla_\theta \log \pi_\theta(a_t | s_t)$ points in the direction that makes action $a_t$ more likely. Multiplying by $G_t$ scales this — if the return was high, push hard in that direction; if low, push the other way.

The Log-Likelihood Trick

How do we differentiate through sampling? We can't backpropagate through np.random.uniform(). The score function estimator (also called the log-likelihood trick or REINFORCE trick) sidesteps this.

For a Bernoulli policy with probability $p$ :

But in practice, we use a simpler form. The gradient of the log-likelihood for a sigmoid output is just $a - p$ :

dlogps.append(action - aprob)  # this IS the score function

This is action - aprob — exactly what the code computes. When action=1 and aprob=0.7, the gradient is 0.3: "make right more likely." This gradient gets backpropagated through the network to update all weights.

Why Reward Standardisation Reduces Variance

REINFORCE is an unbiased estimator of the policy gradient, but it has high variance. A single trajectory might get lucky or unlucky, leading to noisy gradient estimates.

Subtracting a baseline $b$ from the return doesn't change the expected gradient (it's still unbiased) but can dramatically reduce variance:

Our code uses the mean reward as the baseline:

discounted_epr -= np.mean(discounted_epr)

With this baseline, actions that performed better than average get reinforced, and below-average actions get discouraged. Without it, if all rewards are positive (as in CartPole, where every alive step gives +1), the gradient would reinforce all actions — just some more than others. The baseline makes the signal much cleaner.

Hyperparameters

Parameter	What it controls	Value	Original (v0)	Why changed
`H`	Hidden neurons	100	100	Unchanged — enough for CartPole
`learning_rate`	RMSProp step size	1e-3	1e-4	Longer episodes → more samples per update → can use larger steps
`gamma`	Discount factor	0.99	0.95	Longer episodes (500 vs 200) need longer-horizon credit assignment
`batch_size`	Episodes per update	5	5	Unchanged
`decay_rate`	RMSProp memory	0.99	0.99	Unchanged
`epsilon`	RMSProp stability	1e-5	1e-5	Unchanged

The original code targeted CartPole-v0 (max 200 steps) and used a custom "blame window" that penalised only the last N=10 actions. This worked well for short episodes but created a performance ceiling on CartPole-v1 (max 500 steps) — as episodes get longer, the blame signal becomes proportionally smaller and gets washed out after standardisation.

The fix is textbook REINFORCE: standard full-horizon discounting with gamma=0.99. The higher gamma ensures actions 100+ steps before the terminal state still receive meaningful credit. The higher learning rate (1e-3) compensates for the longer episodes providing more gradient samples per update.

Learning rate is critical. Too high and the policy oscillates wildly. Too low and learning crawls. With standard discounting and gamma=0.99, lr=1e-3 converges in ~3,000 episodes.

Value-Based vs Policy-Based: When to Use Each

	DQN (Value-Based)	REINFORCE (Policy-Based)
Outputs	Q-values for each action	Action probabilities
Action selection	Argmax (deterministic)	Sample (stochastic)
Exploration	Epsilon-greedy (bolted on)	Built-in (stochastic policy)
Action spaces	Discrete only	Discrete and continuous
Sample efficiency	Higher (replay buffer)	Lower (on-policy, no replay)
Stability	Needs replay + target net	Needs variance reduction
Convergence	To optimal Q → optimal policy	Directly to (locally) optimal policy

Use policy gradients when:

The action space is continuous (robotics, continuous control)
You want a stochastic policy (exploration, multi-modal strategies)
The policy is simpler than the value function

Use DQN when:

The action space is small and discrete
Sample efficiency matters (you can't run millions of episodes)
You have access to a simulator for experience replay

When NOT to Use Vanilla REINFORCE

High-dimensional action spaces — Variance becomes unmanageable. Use actor-critic methods (A2C, PPO) that learn a value function baseline
Sample-scarce settings — REINFORCE is on-policy: every trajectory is used once then discarded. Off-policy methods like DQN are far more sample-efficient
Long episodes — The credit assignment problem worsens. Which of the 10,000 actions caused success? Actor-critic methods handle this better with per-step value estimates
When stability matters — Vanilla REINFORCE can have large policy swings. PPO clips the gradient to prevent catastrophic updates

Deep Dive: The Papers

Williams (1992) — The REINFORCE Paper

The REINFORCE algorithm was introduced by Ronald J. Williams in his 1992 paper Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning. This is one of the foundational papers in policy gradient methods.

Williams framed the problem precisely: given a stochastic network (he called it a "connectionist network") that produces actions probabilistically, how do we adjust the weights to maximise expected reward?

His key result:

"For each such algorithm, convergence to at least a local maximum in expected reinforcement is assured by a simple condition on the sequence of values used to scale the gradient estimate."

The REINFORCE update rule from the paper:

Where:

$w_{ij}$ — weight from unit $j$ to unit $i$
$\alpha_{ij}$ — learning rate
$r$ — the reinforcement (reward)
$b_{ij}$ — a reinforcement baseline (reduces variance without introducing bias)
$g_i$ — the probability of the output of unit $i$

In our code, this maps directly:

$\frac{\partial \ln g_i}{\partial w_{ij}}$ → dlogps (the score function, computed via backprop)
$r - b_{ij}$ → discounted_epr (standardised, which implicitly subtracts a mean baseline)
The product epdlogp *= discounted_epr is the REINFORCE update

Williams proved that this estimator is unbiased: in expectation, it equals the true policy gradient regardless of the baseline $b$ . The baseline only affects variance — choosing it well (e.g., as the mean reward) can dramatically speed convergence.

Sutton, McAllester, Singh & Mansour (1999) — The Policy Gradient Theorem

The policy gradient theorem, proved in Policy Gradient Methods for Reinforcement Learning with Function Approximation, generalised Williams' result to arbitrary function approximators:

Where $d^{\pi}(s)$ is the stationary state distribution under policy $\pi$ . The theorem shows that the policy gradient doesn't depend on the gradient of the state distribution — a surprising and essential result that makes policy gradient methods practical.

Karpathy (2016) — Deep RL: Pong from Pixels

The original code was inspired by Andrej Karpathy's influential blog post Deep Reinforcement Learning: Pong from Pixels. Karpathy demonstrated that a simple two-layer network trained with REINFORCE could learn to play Atari Pong directly from raw pixels — all in ~130 lines of Python.

The key architectural choices we inherit:

Two-layer network with ReLU hidden layer and sigmoid output
Manual NumPy backprop — no framework dependency
RMSProp as the optimiser
Reward discounting with standardisation

Karpathy's version used the raw Pong pixel difference as input (preprocessed frames). Our CartPole version uses the 4-dimensional state vector directly, but the algorithm is identical.

The REINFORCE Algorithm (Pseudocode)

From Sutton & Barto (2018), Section 13.3:

Initialize policy parameters θ arbitrarily
For each episode:
    Generate trajectory τ = (s₀, a₀, r₁, s₁, a₁, ..., sT) following π_θ
    For each step t = 0, 1, ..., T-1:
        Gₜ ← Σ_{k=t+1}^{T} γ^{k-t-1} rₖ       (return from step t)
        θ ← θ + α γ^t Gₜ ∇_θ ln π_θ(aₜ|sₜ)    (policy gradient update)

Our Implementation vs the Theory

Theory (Williams, 1992)	Our code
$\frac{\partial \ln g_i}{\partial w_{ij}}$ (score function)	`action - aprob` backpropagated through network
$r - b_{ij}$ (reward minus baseline)	`discounted_epr` after mean subtraction
$\Delta w = \alpha (r-b) \nabla \ln \pi$	`epdlogp *= discounted_epr` then `backward()`
Single-sample update	Batch of 5 episodes
Generic optimiser	RMSProp (adaptive learning rates)

What Came After

REINFORCE spawned the entire family of modern policy gradient methods:

Actor-Critic (Konda & Tsitsiklis, 2000) — Learn a value function baseline alongside the policy
A3C (Mnih et al., 2016) — Asynchronous actor-critic with parallel workers
PPO (Schulman et al., 2017) — Clipped surrogate objective for stable updates; the default algorithm behind ChatGPT's RLHF
DDPG (Lillicrap et al., 2016) — Continuous-action policy gradients with a deterministic policy
SAC (Haanoja et al., 2018) — Entropy-regularised policy gradients for robust exploration

Historical Context

Williams (1992) — REINFORCE: the first general-purpose policy gradient algorithm
Sutton et al. (1999) — Policy gradient theorem: proves the gradient formula for function approximators
Kakade (2001) — Natural policy gradients: use Fisher information for better updates
Schulman et al. (2015) — TRPO: trust regions for policy optimisation
Schulman et al. (2017) — PPO: simplified TRPO with clipped objective; now the industry standard
OpenAI (2017) — RLHF: policy gradients for aligning language models

Interactive Tools

Q-Learning Visualiser — Compare value-based RL with the policy gradient approach covered in this post

Deep Q-Networks: Experience Replay and Target Networks — Value-based RL with neural networks, the approach we move beyond here
Q-Learning from Scratch — Tabular value-based RL, the foundation for understanding the value→policy shift
Backpropagation Demystified — We implement backprop manually here too; this post covers the fundamentals
Genetic Algorithms — Gradient-free optimisation, an alternative to policy gradients for non-differentiable objectives

Try It Yourself

The interactive notebook includes exercises:

Remove reward shaping — Replace discount_rewards with standard full-horizon discounting. How does training speed change?
Vary the blame window — Try $N \in \{5, 10, 20, 50\}$ . How does the learning curve respond?
Remove standardisation — Comment out the mean/std normalisation of rewards. Does the agent still learn?
Gamma sweep — Try $\gamma \in \{0.8, 0.95, 0.99, 0.999\}$ . How does the discount factor affect learning?
Compare with DQN — Plot the REINFORCE learning curve alongside DQN's on the same axes. Which learns faster? Which is more stable?

REINFORCE showed that you can optimise a policy directly — no Q-values needed. The gradient signal is noisy, but with variance reduction (baselines, reward standardisation, batch updates) it works remarkably well. The same core idea — gradient ascent on expected reward — underpins every modern policy gradient algorithm, from PPO to the RLHF that aligns large language models. The fundamentals haven't changed since Williams (1992); only the variance reduction has gotten better.

Frequently Asked Questions

What is the difference between policy gradient methods and value-based methods like DQN?

Value-based methods learn a value function (such as Q-values) and derive a policy indirectly by choosing the action with the highest value. Policy gradient methods parameterise the policy directly and optimise it via gradient ascent on expected reward. The key advantage of policy gradients is that they naturally handle continuous action spaces and produce stochastic policies with built-in exploration.

Why is reward standardisation important in REINFORCE?

Without standardisation, all rewards in CartPole are positive (+1 per timestep), so the gradient reinforces every action, just some more than others. Subtracting the mean makes roughly half the actions receive positive reinforcement (above average) and half negative (below average). This acts as a simple baseline that dramatically reduces gradient variance and stabilises training.

What does the discount factor gamma control?

Gamma determines how much weight future rewards receive relative to immediate rewards. A value of 0.99 means an action taken 100 steps before the end still receives about 37% of the terminal reward. Higher gamma values encourage long-term planning but increase variance, while lower values make the agent more short-sighted but produce more stable gradients.

Why was RMSProp chosen as the optimiser instead of plain gradient ascent?

RMSProp maintains a running average of squared gradients and adapts the learning rate for each parameter individually. Parameters with consistently large gradients get smaller effective learning rates, preventing them from dominating the update. This adaptive behaviour is especially important in reinforcement learning, where gradient magnitudes can vary wildly across episodes and parameters.

When should I use REINFORCE versus more advanced algorithms like PPO?

Vanilla REINFORCE is best suited for simple environments, educational purposes, and situations where you want full control over the implementation. For complex environments, long episodes, or production systems, PPO and other actor-critic methods are preferred because they use a learned value function baseline that reduces variance and clip the gradient to prevent catastrophic policy updates.

Why does REINFORCE accumulate gradients over multiple episodes before updating?

Accumulating gradients over a batch of episodes (5 in this implementation) reduces the variance of the gradient estimate. A single episode can be noisy due to the stochastic nature of both the policy and the environment. Averaging over multiple episodes produces a smoother, more reliable gradient signal, similar to minibatch gradient descent in supervised learning.

Deep Q-Networks: Experience Replay and Target Networks

Berkan Sesen — Mon, 06 Apr 2026 09:46:10 +0000

In the Q-learning post, we trained an agent to navigate a 4×4 frozen lake using a simple lookup table — 16 states × 4 actions = 64 numbers. But what happens when the state space isn't a grid?

CartPole has four continuous state variables: cart position, cart velocity, pole angle, and pole angular velocity. Even if you discretised each into 100 bins, you'd need 100⁴ = 100 million Q-values. An Atari game frame is 210×160 pixels with 128 colours — roughly $10^{18}$ possible states. Tables don't work here.

The solution: replace the Q-table with a neural network. Feed in the state, get out Q-values for every action. But naively combining neural networks with Q-learning is unstable — the network chases a moving target while training on correlated sequential data. DeepMind solved both problems with two elegant tricks: experience replay and a target network.

By the end of this post, you'll implement a Deep Q-Network from scratch in PyTorch, train it to balance a pole, and understand why these two tricks were the key insight behind Mnih et al. (2013) — the paper that launched modern deep reinforcement learning.

The Problem: CartPole

OpenAI's CartPole-v1 is a classic control problem: balance a pole on a moving cart by pushing left or right.

          |
          |    ← pole (keep upright!)
          |
    ┌─────┴─────┐
    │    cart    │
    └───────────┘
  ◄── push left    push right ──►
  ─────────────────────────────────

State (4 continuous values):

Cart position ( $x$ )
Cart velocity ( $\dot{x}$ )
Pole angle ( $\theta$ )
Pole angular velocity ( $\dot{\theta}$ )

Actions: Push left (0) or push right (1).

Reward: +1 for every timestep the pole stays upright. The episode ends when the pole tilts beyond ±12° or the cart moves beyond ±2.4 from centre. Maximum score is 500 (the episode is truncated there).

Unlike FrozenLake's 16 discrete states, CartPole's state space is continuous and 4-dimensional. A Q-table is useless here — we need function approximation.

Quick Win: Run the Algorithm

Let's see DQN in action. Click the badge to open the interactive notebook:

The GIF shows the agent's progress: in early episodes the pole topples immediately, but after training with experience replay and a target network, it learns to balance for the full 500 steps.

import numpy as np
import random
from collections import deque

import torch
import torch.nn as nn
import torch.optim as optim
import gymnasium as gym

# --- Neural network: maps state → Q-values for each action ---
class QNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim),
        )

    def forward(self, x):
        return self.net(x)

# --- Experience replay buffer ---
class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)
        return (np.array(states), np.array(actions), np.array(rewards),
                np.array(next_states), np.array(dones))

    def __len__(self):
        return len(self.buffer)

# --- DQN Agent ---
class DQNAgent:
    def __init__(self, state_dim, action_dim, hidden_dim=128, lr=1e-3,
                 gamma=0.99, buffer_size=10000, batch_size=64,
                 target_update_freq=5, epsilon_start=1.0,
                 epsilon_end=0.05, epsilon_step=0.0005, min_buffer=500):
        self.action_dim = action_dim
        self.gamma = gamma
        self.batch_size = batch_size
        self.target_update_freq = target_update_freq
        self.min_buffer = min_buffer

        # Linear epsilon decay — only during training, not during buffer collection.
        # Rate from the original CartPole code: epsilon -= 1/(n_episodes/0.05)
        self.epsilon = epsilon_start
        self.epsilon_end = epsilon_end
        self.epsilon_step = epsilon_step

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.q_network = QNetwork(state_dim, action_dim, hidden_dim).to(self.device)
        self.target_network = QNetwork(state_dim, action_dim, hidden_dim).to(self.device)
        self.target_network.load_state_dict(self.q_network.state_dict())

        self.optimizer = optim.Adam(self.q_network.parameters(), lr=lr)
        self.buffer = ReplayBuffer(buffer_size)

    def select_action(self, state):
        if random.random() < self.epsilon:
            return random.randrange(self.action_dim)
        with torch.no_grad():
            state_t = torch.FloatTensor(state).unsqueeze(0).to(self.device)
            return self.q_network(state_t).argmax(dim=1).item()

    def train_step(self):
        if len(self.buffer) < max(self.batch_size, self.min_buffer):
            return None

        states, actions, rewards, next_states, dones = self.buffer.sample(
            self.batch_size
        )

        states_t = torch.FloatTensor(states).to(self.device)
        actions_t = torch.LongTensor(actions).to(self.device)
        rewards_t = torch.FloatTensor(rewards).to(self.device)
        next_states_t = torch.FloatTensor(next_states).to(self.device)
        dones_t = torch.FloatTensor(dones).to(self.device)

        # Current Q-values: Q(s, a) from the online network
        q_values = self.q_network(states_t).gather(1, actions_t.unsqueeze(1)).squeeze()

        # Target Q-values: r + γ max_a' Q_target(s', a') from the frozen network
        with torch.no_grad():
            next_q_values = self.target_network(next_states_t).max(dim=1).values
            targets = rewards_t + self.gamma * next_q_values * (1 - dones_t)

        loss = nn.MSELoss()(q_values, targets)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        # Linear epsilon decay — only during training
        self.epsilon = max(self.epsilon_end, self.epsilon - self.epsilon_step)

        return loss.item()

    def update_target_network(self):
        self.target_network.load_state_dict(self.q_network.state_dict())

# --- Training loop with early stopping ---
import copy

env = gym.make("CartPole-v1")
agent = DQNAgent(
    state_dim=env.observation_space.shape[0],
    action_dim=env.action_space.n,
)

n_episodes = 500
early_stop_reward = 400      # stop when rolling avg hits this
early_stop_window = 50       # rolling window size
rewards_history = []
best_avg_reward = -float('inf')
best_weights = None

for episode in range(n_episodes):
    state, _ = env.reset()
    total_reward = 0

    for step in range(500):
        action = agent.select_action(state)
        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated

        agent.buffer.push(state, action, reward, next_state, done)
        agent.train_step()

        state = next_state
        total_reward += reward
        if done:
            break

    if episode % agent.target_update_freq == 0:
        agent.update_target_network()

    rewards_history.append(total_reward)

    # Track best weights and check for early stopping
    if len(rewards_history) >= early_stop_window:
        avg = np.mean(rewards_history[-early_stop_window:])
        if avg > best_avg_reward:
            best_avg_reward = avg
            best_weights = copy.deepcopy(agent.q_network.state_dict())
        if avg >= early_stop_reward:
            print(f"Early stopping at episode {episode} (avg reward: {avg:.1f})")
            break

    if episode % 50 == 0:
        avg = np.mean(rewards_history[-50:])
        print(f"Episode {episode:3d} | Avg reward: {avg:6.1f} | Epsilon: {agent.epsilon:.3f}")

# Restore best weights
if best_weights is not None:
    agent.q_network.load_state_dict(best_weights)

env.close()
print(f"\nFinal avg reward (last 50): {np.mean(rewards_history[-50:]):.1f}")

The result: The agent learns to balance the pole and early stopping kicks in once the 50-episode rolling average hits 400. We also save the best weights, because DQN can suffer from catastrophic forgetting — performance collapses if you keep training past the sweet spot. Compare this to tabular Q-learning: here we handle a continuous state space with just a small neural network instead of an impossibly large table.

Visualise the Learning

import matplotlib.pyplot as plt

rolling = np.convolve(rewards_history, np.ones(20)/20, mode='valid')

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(rolling, 'b-', linewidth=0.8)
ax.axhline(y=500, color='g', linestyle='--', alpha=0.5, label='Max score (500)')
ax.set_xlabel('Episode')
ax.set_ylabel('Reward (20-episode rolling avg)')
ax.set_title('DQN on CartPole-v1')
ax.set_ylim(0, 550)
ax.legend()
ax.grid(True, alpha=0.3)
plt.show()

The reward curve shows two phases: the agent plateaus around 200 while the replay buffer fills with diverse experiences, then rapidly improves once it starts exploiting what it's learned. Early stopping halts training once the rolling average hits 400 — without it, continued training often leads to catastrophic forgetting where performance suddenly collapses.

What Just Happened?

We replaced the Q-table with a neural network, but two critical tricks made it work. Without them, training is unstable or fails entirely.

Trick 1: Experience Replay

In tabular Q-learning, we updated $Q(s, a)$ immediately after each transition. With a neural network, this is catastrophic. Consecutive experiences are highly correlated — if the cart drifts left for 10 steps, the network sees 10 similar "going left" transitions in a row and overfits to that pattern.

Experience replay breaks this correlation:

class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        # ...

Every transition $(s, a, r, s', \text{done})$ is stored in a fixed-size buffer. During training, we sample a random minibatch from the buffer — not the latest transitions, but a mix of old and new experiences from different parts of the state space.

Think of it like studying for an exam: instead of re-reading the last chapter (correlated, recent data), you shuffle all your flash cards and quiz yourself on a random mix (uncorrelated, diverse data).

The original code used exactly this pattern — a deque with a fixed capacity:

self.experience_pool = deque([], pool_size)  # from the original dqn.py

Mnih et al. used a buffer of 1 million transitions for Atari. Our CartPole agent uses 10,000 — more than enough for this simpler problem.

Trick 2: Target Network (Frozen Q-hat)

The second instability comes from the training target itself. In Q-learning, we update toward:

But the Q-values on the right come from the same network we're updating. It's like a dog chasing its tail — every gradient step shifts the target, causing oscillations or divergence.

The fix: maintain a separate target network $\hat{Q}$ (called frozen_model in the original code, target_network in ours). This network is a periodic copy of the online network:

# Online network: updated every step
q_values = self.q_network(states_t).gather(1, actions_t.unsqueeze(1))

# Target network: frozen copy, updated every N episodes
with torch.no_grad():
    next_q_values = self.target_network(next_states_t).max(dim=1).values
    targets = rewards_t + self.gamma * next_q_values * (1 - dones_t)

The target network is frozen between updates. Every target_update_freq episodes, we copy the online network's weights into it:

if episode % agent.target_update_freq == 0:
    agent.update_target_network()

This gives the online network a stable target to train against. The original code explored different freeze intervals:

freeze_everys = [1, 5, 10]  # from the original cartpole_q_net.py

Updating every episode (freeze_every=1) essentially disables the trick. Freezing for longer (5 or 10 episodes) gives more stable training but slightly slower adaptation.

The Training Step

Putting it together, each training step:

Sample a random minibatch from the replay buffer
Compute current Q-values from the online network: $Q(s, a)$
Compute target Q-values from the frozen network: $r + \gamma \max_{a'} \hat{Q}(s', a')$
Minimise MSE loss between current and target: $\mathcal{L} = (Q(s,a) - y)^2$
Backpropagate and update only the online network's weights

loss = nn.MSELoss()(q_values, targets)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()

This is supervised learning with self-generated labels. The "label" for $Q(s, a)$ is the bootstrapped estimate $r + \gamma \max_{a'} \hat{Q}(s', a')$ — not a ground truth, but a better estimate than what we had.

The Full Picture

┌──────────────┐   action   ┌─────────────┐
│  Environment ├───────────►│   Agent     │
│  (CartPole)  │◄───────────┤  ε-greedy   │
└──────┬───────┘  s,r,s',d  └──────┬──────┘
       │                           │
       │  (s,a,r,s',done)         │
       ▼                           │
┌──────────────┐                   │
│   Replay     │  random batch     │
│   Buffer     ├───────────────────┘
│  (10,000)    │         │
└──────────────┘         ▼
                  ┌──────────────┐     copy weights
                  │  Online Net  ├────────────────►┌──────────────┐
                  │   Q(s,a)     │  every N eps     │  Target Net  │
                  └──────┬───────┘                  │  Q̂(s',a')    │
                         │                          └──────┬───────┘
                         │    MSE loss                     │
                         └─────────── vs ──────────────────┘
                              Q(s,a) ≈ r + γ max Q̂(s',a')

Going Deeper

Why Naive Q-Learning with Neural Nets Fails

There are three reasons combining neural networks with Q-learning is unstable — Sutton & Barto call it the deadly triad:

Correlated samples — Sequential experience creates biased gradient estimates. A batch of "cart drifting left" transitions teaches the network to handle left-drifts but destroys its knowledge of right-drifts. Experience replay fixes this.
Non-stationary targets — The target $r + \gamma \max_{a'} Q(s', a')$ changes with every gradient step because $Q$ is the network we're updating. The target network fixes this by providing a frozen target that only changes periodically.
Function approximation — A neural network generalises: updating Q-values for one state affects nearby states. This can create feedback loops where overestimation in one region cascades. This is partially addressed by the target network and is further improved by Double DQN, which decouples action selection from value estimation.

Practical Tip: Early Stopping

Experience replay and the target network reduce instability, but they don't eliminate it. In practice, DQN agents often exhibit catastrophic forgetting — performance peaks, then collapses as continued training destabilises the network. This is visible in the ablation study: even the full DQN's reward curve can crash if training runs too long.

The fix is simple and borrowed from supervised learning: monitor a rolling average of the reward and stop training once it plateaus. Save the best weights along the way:

best_avg_reward = -float('inf')
best_weights = None

for episode in range(n_episodes):
    # ... training loop ...

    if len(rewards_history) >= 50:
        avg = np.mean(rewards_history[-50:])
        if avg > best_avg_reward:
            best_avg_reward = avg
            best_weights = copy.deepcopy(agent.q_network.state_dict())
        if avg >= 400:  # stop when converged
            break

# Restore best weights
agent.q_network.load_state_dict(best_weights)

This pattern applies broadly: any RL agent trained with function approximation can overfit or destabilise. Always track the best checkpoint.

Hyperparameters

Parameter	What it controls	Our value	DeepMind Atari
`hidden_dim`	Network capacity	128	~3 conv layers + FC
`lr`	Learning rate	1e-3	2.5e-4 (RMSprop)
`gamma`	Discount factor	0.99	0.99
`buffer_size`	Replay buffer capacity	10,000	1,000,000
`min_buffer`	Fill before training	500	50,000
`batch_size`	Minibatch size	64	32
`target_update_freq`	Target network sync	Every 5 episodes	Every 10,000 steps
`epsilon_start`	Initial exploration	1.0	1.0
`epsilon_end`	Final exploration	0.05	0.1
`epsilon_step`	Exploration decay rate	0.0005 per training step	Linear over 1M frames

Buffer size matters. The original code experimented with pool sizes of 500 vs 1,500 and found the larger buffer consistently improved learning. DeepMind used 1 million transitions — enough to hold roughly the last ~4 days of gameplay.

Batch size matters. The original code tested batch sizes of 16, 32, 64, and 128, finding that "increasing from 16 → 32 → 64 → 128 consistently boosts agent's training speed." Larger batches give lower-variance gradient estimates, but eventually hit diminishing returns.

Gamma requires care. The original code's author noted: "Initially I was using $\gamma = 0.75$ for the CartPole problem which is grossly small." With $\gamma = 0.75$ , a reward 10 steps away is worth $0.75^{10} \approx 0.06$ — the agent becomes too myopic to learn long-horizon balancing. With $\gamma = 0.99$ , that same reward is worth $0.99^{10} \approx 0.90$ .

Epsilon Annealing

The original code contained a hard-won insight about epsilon decay:

# The rate with which you anneal epsilon is very crucial in the training quality.
# We found that we were annealing too early, not allowing the agent to explore
# sufficiently.

Epsilon controls the explore-exploit trade-off — just like in tabular Q-learning. But with DQN, the stakes are higher: the network needs diverse training data to generalise well. Anneal too fast and the agent gets stuck in a local optimum; too slow and training wastes time on random actions.

Two critical details from the original code:

Fill the buffer first — epsilon stays at 1.0 (pure random) until min_buffer transitions are collected. This seeds the replay buffer with diverse experiences before any training begins.
Linear decay during training only — epsilon decrements by a fixed step (0.0005) each training step, decaying from 1.0 to 0.05 over ~1,900 training steps. DeepMind used a similar linear decay from 1.0 to 0.1 over 1 million frames.

From CartPole to Atari

Our CartPole DQN takes a 4-dimensional state vector as input. DeepMind's Atari DQN takes raw pixels: 84×84 grayscale frames, stacked 4 deep (to capture motion). The key architectural difference is convolutional layers for spatial feature extraction:

Atari DQN Architecture:
  Input: 84×84×4 (4 stacked frames)
  → Conv2D(32, 8×8, stride 4) → ReLU
  → Conv2D(64, 4×4, stride 2) → ReLU
  → Conv2D(64, 3×3, stride 1) → ReLU
  → Flatten → Dense(512) → ReLU
  → Dense(num_actions)          # one output per action

The same two tricks — experience replay and target network — are what made this work. The network architecture is straightforward; the training stability innovations are the real contribution.

When NOT to Use DQN

Continuous action spaces — DQN requires $\arg\max_a Q(s, a)$ over discrete actions. For continuous control (robotic arms, throttle), use policy gradient methods like DDPG or SAC
Very long episodes — DQN's bootstrapping can propagate errors over long horizons. For episodic tasks with clear terminal rewards, Monte Carlo methods can be more stable
Simple environments — If the state space is small and discrete, tabular Q-learning is simpler, interpretable, and guaranteed to converge. Don't use a neural network when a table will do
Multi-agent settings — DQN assumes a stationary environment. Other agents make the environment non-stationary, breaking core assumptions

Ablation: What Happens Without the Tricks?

To see why experience replay matters, we train a variant without it — using a tiny 64-transition buffer that can't break temporal correlations:

# Full DQN with experience replay (early stopped)
agent_full = DQNAgent(state_dim=4, action_dim=2)

# No experience replay: tiny buffer, small batch
agent_no_replay = DQNAgent(state_dim=4, action_dim=2, buffer_size=64, batch_size=8, min_buffer=8)

The contrast is stark:

Full DQN (blue) — Converges to ~400+ reward. Early stopping captures the peak before catastrophic forgetting can set in
No experience replay (red) — Stuck below 150 for 500 episodes. Without random sampling from a large buffer, the network trains on correlated sequential data and never generalises

This matches what DeepMind found: experience replay was the critical innovation. The correlated, non-i.i.d. nature of sequential RL data is fundamentally incompatible with stable neural network training — random minibatches from a large buffer fix this.

Deep Dive: The Papers

Mnih et al. (2013) — The NIPS Workshop Paper

The story begins at DeepMind in 2013. Volodymyr Mnih and colleagues published Playing Atari with Deep Reinforcement Learning at the NIPS Deep Learning Workshop. The goal was ambitious: learn to play Atari 2600 games directly from pixels, using a single architecture and hyperparameter set across all games.

The key quote from the paper:

"We demonstrate that a convolutional neural network can learn successful control policies from raw video data [...] using only the reward signal. [...] The network was not provided with any game-specific information or hand-designed visual features."

Previous attempts at combining neural networks with Q-learning had failed. The paper identified two reasons:

"Firstly, the sequence of observations in RL is correlated, unlike the independent and identically distributed (i.i.d.) assumption of most deep learning methods. Secondly, the data distribution changes as the agent learns new behaviours."

Their solution was experience replay, borrowed from Lin (1992):

"We store the agent's experiences at each time step, $e_t = (s_t, a_t, r_t, s_{t+1})$ , in a data set $\mathcal{D} = e_1, \ldots, e_N$ , pooled over many episodes into a replay memory."

Mnih et al. (2015) — The Nature Paper

Two years later, the expanded version appeared in Nature as Human-level control through deep reinforcement learning. The key addition was the target network:

"The second modification [...] is to use a separate network for generating the targets $y_j$ in the Q-learning update. More precisely, every $C$ updates we clone the network $Q$ to obtain a target network $\hat{Q}$ and use $\hat{Q}$ for generating the Q-learning targets for the following $C$ updates."

The loss function, from the paper:

Where:

$\theta_i$ — weights of the online network at iteration $i$
$\theta_i^{-}$ — weights of the target network (frozen, periodically copied from $\theta$ )
$U(\mathcal{D})$ — uniform random sampling from the replay buffer

The Nature paper tested on 49 Atari games. DQN achieved human-level performance or better on 29 of them — a remarkable result for a single, general-purpose algorithm.

The DQN Algorithm (Paper Pseudocode)

From Algorithm 1 of the Nature paper:

Initialize replay memory D with capacity N
Initialize action-value function Q with random weights θ
Initialize target action-value function Q̂ with weights θ⁻ = θ

For episode = 1, M do:
    Initialise state s₁
    For t = 1, T do:
        With probability ε select random action aₜ
        Otherwise select aₜ = argmax_a Q(sₜ, a; θ)
        Execute action aₜ, observe reward rₜ and next state sₜ₊₁
        Store transition (sₜ, aₜ, rₜ, sₜ₊₁) in D
        Sample random minibatch of transitions (sⱼ, aⱼ, rⱼ, sⱼ₊₁) from D
        Set yⱼ = rⱼ                           if episode terminates at step j+1
             yⱼ = rⱼ + γ max_a' Q̂(sⱼ₊₁, a'; θ⁻)   otherwise
        Perform gradient descent on (yⱼ - Q(sⱼ, aⱼ; θ))² w.r.t. θ
        Every C steps: Q̂ ← Q  (reset target network)

Our Implementation vs the Paper

Paper (DQN, 2015)	Our code
Input: 84×84×4 pixel frames	Input: 4-dim state vector
3 conv layers + 2 FC layers	2 hidden FC layers (128 units)
Replay buffer $\mathcal{D}$ , capacity 1M	`ReplayBuffer`, capacity 10,000
Uniform random sample $U(\mathcal{D})$	`random.sample(self.buffer, batch_size)`
Target network $\hat{Q}(\theta^-)$	`self.target_network`
Clone every $C = 10{,}000$ steps	`update_target_network()` every 10 episodes
RMSprop, $lr = 0.00025$	Adam, $lr = 0.001$
Loss: $(y - Q(s,a;\theta))^2$	`nn.MSELoss()(q_values, targets)`
49 Atari games	CartPole-v1

The architecture differs, but the algorithm is identical. The same two tricks — experience replay and target network — stabilise training whether the input is raw pixels or a 4-dimensional state vector.

What Came After

DQN spawned a wave of improvements:

Double DQN (van Hasselt et al., 2016) — Fixes Q-value overestimation by using the online network to select actions but the target network to evaluate them
Prioritised Experience Replay (Schaul et al., 2016) — Sample transitions with high TD error more often (learn from surprises)
Dueling DQN (Wang et al., 2016) — Separate streams for state value $V(s)$ and advantage $A(s, a)$
Rainbow (Hessel et al., 2018) — Combines all improvements into one agent
Distributional RL (Bellemare et al., 2017) — Learn the full distribution of returns, not just the mean

Historical Context

Watkins (1989) — Q-learning: model-free, off-policy control
Lin (1992) — Experience replay for RL (first proposed)
Riedmiller (2005) — Neural Fitted Q-Iteration (NFQ) — early neural network Q-learning
Mnih et al. (2013) — DQN: experience replay + deep learning on Atari
Mnih et al. (2015) — DQN + target network in Nature, human-level on 29/49 games
Silver et al. (2016) — AlphaGo: deep RL defeats world Go champion

Interactive Tools

Q-Learning Visualiser — Watch tabular Q-learning train step-by-step on grid worlds before scaling up to DQN

Q-Learning from Scratch — Tabular Q-learning, the foundation DQN builds on
Backpropagation Demystified — The gradient engine powering DQN's neural network
Genetic Algorithms — Evolution-based optimisation, an alternative to gradient-based RL

Try It Yourself

The interactive notebook includes exercises:

Ablation study — Train with target_update_freq=1 (no target network) and buffer_size=1 (no replay). How do the reward curves compare to full DQN?
Gamma sweep — Try $\gamma \in \{0.5, 0.9, 0.99, 0.999\}$ . How does the discount factor affect convergence speed?
Network size — Try hidden_dim of 32, 64, 128, and 256. Is bigger always better for CartPole?
Soft target updates — Replace the hard copy with Polyak averaging: $\theta^- \leftarrow \tau \theta + (1-\tau) \theta^-$ with $\tau = 0.005$ . Does this improve stability?
Double DQN — Modify the target to use the online network for action selection: $a^* = \arg\max_{a'} Q(s', a'; \theta)$ , then evaluate with the target network: $\hat{Q}(s', a^*; \theta^-)$ . Does this reduce overestimation?

DQN showed that deep learning and reinforcement learning could be combined — but only with the right tricks. Experience replay breaks temporal correlations, the target network provides stable training signals, and together they turned an unstable combination into a system that matched human performance on Atari. Next up: policy gradient methods, which skip Q-values entirely and directly optimise the policy.

Frequently Asked Questions

What is a Deep Q-Network and how does it differ from tabular Q-learning?

A DQN replaces the Q-table with a neural network that takes a state as input and outputs Q-values for all actions. This allows it to handle environments with large or continuous state spaces where a table would be impossibly large. The core Q-learning update remains the same, but the function approximation introduces new stability challenges.

Why is experience replay important for DQN?

Without experience replay, the network trains on consecutive, highly correlated transitions, which violates the i.i.d. assumption of gradient descent and causes unstable learning. Experience replay stores transitions in a buffer and samples random mini-batches for training, breaking correlations and allowing each experience to be reused multiple times.

What is a target network and why is it needed?

In Q-learning, the same network both selects actions and computes target values, creating a moving target problem that destabilises training. A target network is a separate, slowly updated copy used only to compute targets. This decouples the target from the current network's rapidly changing weights, stabilising learning significantly.

How do I choose the replay buffer size?

A buffer that is too small loses useful past experiences and reduces diversity. A buffer that is too large wastes memory and may contain outdated transitions from a much weaker policy. For most environments, 100,000 to 1,000,000 transitions works well. Start with 100,000 and increase if training is unstable.

What is the difference between DQN and Double DQN?

Standard DQN tends to overestimate Q-values because it uses the max operator to both select and evaluate actions. Double DQN fixes this by using the online network to select the best action but the target network to evaluate it. This simple change often improves performance significantly with no computational overhead.

Q-Learning from Scratch: Navigating the Frozen Lake

Berkan Sesen — Sat, 04 Apr 2026 08:24:35 +0000

Imagine you're standing on a frozen lake. Your goal is on the far side, but there are holes in the ice — fall in and it's game over. Worse, the ice is slippery: when you try to go right, you might slide up or down instead. You have no map, no instructions. All you can do is try, fail, and gradually learn which moves lead to safety.

This is exactly what Q-learning solves. The agent learns a value for every state-action pair — "how good is it to take action A from state S?" — purely from trial and error. No model of the environment needed, no supervision, just rewards.

By the end of this post, you'll implement Q-learning from scratch, train an agent to navigate OpenAI's FrozenLake environment, and understand the Bellman equation that makes it all work. You'll also see why exploration matters — and what happens when an agent gets greedy too early.

The algorithm was introduced by Watkins (1989) in his PhD thesis and its convergence was proven by Watkins & Dayan (1992).

The Problem: FrozenLake

OpenAI's FrozenLake is a 4×4 grid world:

SFFF    S = Start
FHFH    F = Frozen (safe)
FFFH    H = Hole (game over)
HFFG    G = Goal (reward = 1)

The agent can move Left, Down, Right, Up. But on slippery ice, each action has only a 1/3 chance of going in the intended direction — there's a 1/3 chance of sliding in each perpendicular direction. The only reward is +1 for reaching the Goal. Fall in a Hole or run out of steps and you get 0.

This is a model-free problem: the agent doesn't know the transition probabilities or the reward structure. It must learn entirely from experience.

Quick Win: Run the Algorithm

Let's see Q-learning in action. Click the badge to open the interactive notebook:

The GIF shows the agent's learned policy evolving from random arrows (no idea what to do) to a coherent strategy that avoids holes and reaches the goal. The colour intensity shows $V(s) = \max_a Q(s, a)$ — how valuable each state is.

import numpy as np
import gymnasium as gym

def q_learning(env, n_episodes=20000, alpha=0.1, gamma=0.99,
               epsilon_start=1.0, epsilon_end=0.01, decay_rate=1e-3):
    """Q-learning with epsilon-greedy exploration."""
    Q = np.zeros((env.observation_space.n, env.action_space.n))
    rewards = []

    epsilon = epsilon_start
    for episode in range(n_episodes):
        state, _ = env.reset()
        total_reward = 0

        for step in range(100):
            # Epsilon-greedy action selection
            if np.random.random() < epsilon:
                action = env.action_space.sample()
            else:
                action = np.argmax(Q[state])

            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated

            # Q-learning update (Bellman equation)
            Q[state, action] += alpha * (
                reward + gamma * np.max(Q[next_state]) - Q[state, action]
            )

            total_reward += reward
            state = next_state
            if done:
                break

        rewards.append(total_reward)
        epsilon = epsilon_end + (epsilon_start - epsilon_end) * np.exp(
            -decay_rate * episode
        )

    return Q, rewards

# Train the agent
env = gym.make('FrozenLake-v1', is_slippery=True)
Q, rewards = q_learning(env)

# Show the learned policy
actions = ['←', '↓', '→', '↑']
policy = np.array([actions[a] for a in np.argmax(Q, axis=1)]).reshape(4, 4)
print("Learned policy:")
print(policy)
print(f"\nSuccess rate (last 1000): {np.mean(rewards[-1000:]):.1%}")

The result: After 20,000 episodes, the agent learns a policy that reaches the goal ~67% of the time. Even with a good policy, the slippery ice inevitably pushes you into holes — the theoretical maximum for this environment is about 74%.

Visualise the Learning

import matplotlib.pyplot as plt

rolling = np.convolve(rewards, np.ones(100)/100, mode='valid')

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(rolling, 'b-', linewidth=0.8)
ax.set_xlabel('Episode')
ax.set_ylabel('Success Rate (100-episode rolling avg)')
ax.set_title('Q-Learning on FrozenLake')
ax.set_ylim(0, 1)
ax.grid(True, alpha=0.3)
plt.show()

The reward curve climbs from 0% (random) to ~67% (near-optimal). The initial flat period is pure exploration — the agent is collecting experience before it starts exploiting what it's learned.

What Just Happened?

Q-learning maintains a Q-table — a lookup table mapping every (state, action) pair to a value. For FrozenLake, that's 16 states × 4 actions = 64 entries, all starting at zero.

The Bellman Equation

The core update rule is one line of code:

Q[state, action] += alpha * (
    reward + gamma * np.max(Q[next_state]) - Q[state, action]
)

This implements the temporal difference (TD) update:

Let's break this apart:

$Q(s, a)$ — current estimate of "how good is action a in state s?"
$r$ — immediate reward received
$\gamma \max_{a'} Q(s', a')$ — discounted value of the best action in the next state
$r + \gamma \max_{a'} Q(s', a')$ — the TD target (what we think the true value is)
$\text{target} - Q(s, a)$ — the TD error (how wrong we were)
$\alpha$ — learning rate (how much to correct)

Think of it like updating a GPS estimate: you have a current estimate of the distance ( $Q(s,a)$ ), you observe a new measurement (the target), and you adjust by a fraction ( $\alpha$ ) of the difference.

Epsilon-Greedy Exploration

if np.random.random() < epsilon:
    action = env.action_space.sample()   # explore: random action
else:
    action = np.argmax(Q[state])         # exploit: best known action

With probability $\epsilon$ , the agent takes a random action (exploration). Otherwise, it takes the best action according to its Q-table (exploitation).

Early in training, $\epsilon = 1.0$ — pure exploration. Over time, it decays exponentially:

This mirrors the exploration-exploitation trade-off in MCMC — the Metropolis-Hastings algorithm also balances proposing random moves with accepting good ones.

Why "Off-Policy"?

Q-learning is off-policy: it learns about the greedy policy ( $\max_{a'} Q(s', a')$ ) while following an exploratory policy (epsilon-greedy). This separation is powerful — the agent can explore freely without corrupting what it's learning about the optimal policy.

The alternative, SARSA (on-policy), updates using the action the agent actually took next, not the best possible action. SARSA is more conservative — it learns a safer policy that accounts for its own exploration.

Going Deeper

Hyperparameters

Parameter	What it controls	FrozenLake value	Guidance
`alpha`	Learning rate	0.1	Low for stochastic envs, higher for deterministic
`gamma`	Discount factor	0.99	Close to 1 for long-horizon tasks
`epsilon_start`	Initial exploration	1.0	Always start fully random
`epsilon_end`	Final exploration	0.01	Small but non-zero
`decay_rate`	Exploration decay	0.001	Slower = more exploration
`n_episodes`	Training budget	20,000	Until convergence

gamma is critical. It controls how much the agent values future rewards vs immediate ones:

$\gamma = 0$ — only cares about immediate reward (myopic)
$\gamma = 1$ — values all future rewards equally (farsighted)
$\gamma = 0.95$ — reward 10 steps away is worth $0.95^{10} \approx 0.60$ of its face value

For FrozenLake, the goal is ~6 steps away but stochastic transitions make paths longer. With $\gamma = 0.99$ , a reward 10 steps away is worth $0.99^{10} \approx 0.90$ — the agent strongly values the distant goal.

Slippery vs Non-Slippery

# Deterministic: agent goes exactly where intended
env_easy = gym.make('FrozenLake-v1', is_slippery=False)
Q_easy, rewards_easy = q_learning(env_easy)
print(f"Non-slippery success rate: {np.mean(rewards_easy[-1000:]):.1%}")  # ~100%

# Stochastic: 1/3 chance of sliding sideways
env_hard = gym.make('FrozenLake-v1', is_slippery=True)
Q_hard, rewards_hard = q_learning(env_hard)
print(f"Slippery success rate: {np.mean(rewards_hard[-1000:]):.1%}")  # ~75%

Without slippery ice, the agent achieves ~100% success — the problem becomes trivial. The stochastic transitions are what make FrozenLake interesting: the agent must learn a policy that's robust to noise.

When NOT to Use Tabular Q-Learning

Large state spaces — A 100×100 grid has 10,000 states × 4 actions = 40,000 Q-values. An Atari game has ~10¹⁸ possible states. Tables don't scale — use neural network function approximation instead (Deep Q-Networks)
Continuous states — Temperature, position, velocity can't be tabulated. Use function approximation or discretise
Continuous actions — Q-learning requires $\max_a Q(s,a)$ , which is trivial for discrete actions but requires optimisation for continuous ones. Use policy gradient methods instead

Alternatives:

Deep Q-Networks (DQN) — Replace the Q-table with a neural network
SARSA — On-policy variant, learns a safer policy
Policy gradient methods — Directly optimise the policy, handle continuous actions
Monte Carlo methods — Wait until the end of the episode to update (no bootstrapping)

Deep Dive: The Papers

Watkins (1989)

Christopher Watkins introduced Q-learning in his PhD thesis Learning from Delayed Rewards at Cambridge. The key problem he addressed: how can an agent learn optimal behaviour when rewards are delayed — you don't know if a move was good until many steps later?

His solution was to learn the action-value function $Q(s, a)$ directly, rather than the state-value function $V(s)$ . The crucial insight: by learning $Q$ values, the agent can select optimal actions without needing a model of the environment — it just picks $\arg\max_a Q(s, a)$ .

Convergence Guarantee

Watkins & Dayan proved in their 1992 paper that Q-learning converges to the optimal Q-function $Q^*$ under two conditions:

Every state-action pair is visited infinitely often
The learning rate decays appropriately (but not too fast)

Formally, given the update:

Convergence requires:

The first condition ensures sufficient learning; the second ensures the updates eventually settle. In practice, a constant learning rate works well for tabular problems.

The Bellman Optimality Equation

Q-learning is grounded in Bellman's Dynamic Programming (1957). The optimal Q-function satisfies the Bellman optimality equation:

This is a fixed-point equation: the optimal Q-values are self-consistent. Q-learning iteratively approximates this fixed point using sampled transitions instead of computing the full expectation.

Our Implementation vs the Theory

Theory	Our code
State space $\mathcal{S}$	16 grid cells (0-15)
Action space $\mathcal{A}$	4 directions (L, D, R, U)
Q-function $Q(s,a)$	`Q` numpy array (16×4)
Policy $\pi(s) = \arg\max_a Q(s,a)$	`np.argmax(Q[state])`
Bellman update	`Q[s,a] += alpha * (r + gamma * max(Q[s']) - Q[s,a])`
Exploration $\epsilon$ -greedy	`if random() < epsilon: sample()`

Historical Context

Bellman (1957) — Dynamic programming and the Bellman equation
Samuel (1959) — Checkers-playing program, early TD learning
Sutton (1988) — TD(λ) learning framework
Watkins (1989) — Q-learning: model-free off-policy control
Watkins & Dayan (1992) — Convergence proof for Q-learning
Mnih et al. (2013) — Deep Q-Networks (DQN) — Q-learning with neural nets, playing Atari
Sutton & Barto (2018) — Reinforcement Learning: An Introduction — the definitive textbook

Interactive Tools

Q-Learning Visualiser — Watch Q-learning train step-by-step on grid worlds in the browser

Genetic Algorithms — evolution-based optimisation (no gradient, no model — like Q-learning)
Backpropagation Demystified — the gradient engine that powers Deep Q-Networks
MCMC Island Hopping — another algorithm balancing exploration and exploitation

Try It Yourself

The interactive notebook includes exercises:

Slippery vs non-slippery — Compare convergence curves for is_slippery=True and False. How does the optimal success rate differ?
Gamma sweep — Try $\gamma \in \{0.5, 0.8, 0.95, 0.99\}$ . How does the discount factor affect learning speed and final performance?
SARSA comparison — Modify the update to use the actual next action instead of $\max_{a'}$ . How does the learned policy differ?
8×8 FrozenLake — Switch to FrozenLake8x8-v1 (64 states). Does Q-learning still converge? How many more episodes does it need?
Visualise exploration — Track how many times each state is visited during training. Which states are under-explored?

Understanding Q-learning opens the door to Deep Q-Networks, policy gradients, and modern RL. The core insight — learning action values from delayed rewards through bootstrapping — is the foundation of all value-based reinforcement learning.

Frequently Asked Questions

What is the difference between Q-learning and SARSA?

Q-learning is off-policy: it updates using the best possible next action regardless of what the agent actually does. SARSA is on-policy: it updates using the action the agent actually took next, which means exploration noise feeds into the value estimates. In practice, SARSA learns a more conservative policy that accounts for the agent's own exploratory behaviour.

Why does the Q-table start with all zeros?

Initialising all Q-values to zero is a simple and common choice for environments where rewards are non-negative. Since the only reward in FrozenLake is +1 at the goal, any discovered path to the goal will produce Q-values above zero, naturally encouraging the agent to pursue it. Optimistic initialisation (starting with high values) is an alternative that can speed up exploration.

Can Q-learning handle continuous state spaces?

Not in its tabular form. With continuous states like position or velocity, you cannot maintain a lookup table for every possible state. The standard solution is to replace the Q-table with a function approximator such as a neural network, which is exactly what Deep Q-Networks (DQN) do. Alternatively, you can discretise the state space, though this scales poorly to high dimensions.

Why does the agent only reach ~67-74% success on slippery FrozenLake?

The slippery ice introduces genuine stochasticity that no policy can overcome. Each action has only a 1/3 chance of going in the intended direction, so even with a perfect policy, the agent will sometimes slide into holes. The theoretical maximum success rate for this environment is approximately 74%, meaning the agent's learned policy is near-optimal.

How does the discount factor gamma affect learning?

Gamma controls how much the agent values future rewards relative to immediate ones. A low gamma (e.g. 0.5) makes the agent short-sighted, focusing only on nearby rewards. A high gamma (e.g. 0.99) makes the agent plan further ahead, which is essential in FrozenLake where the goal is several steps away. Setting gamma too close to 1.0 can slow convergence because the agent gives nearly equal weight to very distant outcomes.

Does Q-learning guarantee finding the optimal policy?

Yes, under theoretical conditions: every state-action pair must be visited infinitely often, and the learning rate must decay appropriately. In practice, with a fixed learning rate and epsilon-greedy exploration, Q-learning converges reliably on small tabular problems like FrozenLake. For larger or continuous problems, convergence guarantees no longer hold, and techniques like experience replay and target networks become necessary.

Genetic Algorithms: From Line Fitting to the Travelling Salesman

Berkan Sesen — Fri, 03 Apr 2026 15:14:53 +0000

Imagine you're planning a road trip through 25 cities. The number of possible routes is 25!/2 — roughly 7.8 × 10²⁴, more than the number of stars in the observable universe. You can't try them all. And unlike fitting a line with gradient descent, there's no gradient to follow — the search space is discrete and rugged.

Nature solved this class of problem billions of years ago. Evolution doesn't need gradients. It works through variation and selection: generate diverse candidates, keep the ones that work, combine their traits, add random mutations, and repeat. Genetic algorithms (GAs) formalise this process into a general-purpose optimiser.

By the end of this post, you'll implement GAs from scratch for two problems: fitting a regression line (where we can verify the answer) and the Travelling Salesman Problem (where we can't). You'll understand the core operators — selection, crossover, and mutation — and see how a single evolutionary loop solves fundamentally different problems.

This tutorial draws on Holland's foundational Adaptation in Natural and Artificial Systems (1975) and Whitley's excellent A Genetic Algorithm Tutorial (1994).

Quick Win: Evolving a Line

Let's see evolution in action on a simple problem. We have 50 data points generated from $y = 3x + 8$ (with some noise), and we want to evolve the coefficients $(a, b)$ that best fit the data — no calculus required.

Click the badge to open the interactive notebook:

The GA starts with 100 random $(a, b)$ pairs — wild guesses from $[-10, 10]$ . Each generation, the best-fitting individuals survive and breed. Within ~20 generations, the population converges to the true parameters.

import numpy as np

def ga_line_fit(x, y, pop_size=100, n_elites=20, mutation_rate=0.1,
                generations=50):
    """Evolve (a, b) to minimise RMSE for y = ax + b."""
    pop = np.random.uniform(-10, 10, size=(pop_size, 2))
    history = []

    for gen in range(generations):
        # FITNESS: how well each individual fits the data
        preds = pop[:, 0:1] * x + pop[:, 1:2]
        rmse = np.sqrt(((preds - y) ** 2).mean(axis=1))
        fitness = 1.0 / (rmse + 1e-10)

        best_idx = np.argmax(fitness)
        history.append((pop[best_idx, 0], pop[best_idx, 1], rmse[best_idx]))

        # SELECTION: keep elite performers + sample by fitness
        ranked = np.argsort(-fitness)
        elites = pop[ranked[:n_elites]]
        probs = fitness / fitness.sum()
        parents = pop[np.random.choice(pop_size, pop_size - n_elites, p=probs)]

        # CROSSOVER + MUTATION: average parents, add Gaussian noise
        children = np.empty_like(parents)
        for i in range(len(parents)):
            child = (parents[i] + parents[-(i + 1)]) / 2
            child *= 1 + mutation_rate * np.random.randn(2)
            children[i] = child

        pop = np.vstack([elites, children])

    return history

# Generate noisy data: y = 3x + 8
np.random.seed(42)
x = np.random.rand(50)
y = 3 * x + 8 + 0.3 * np.random.randn(50)

history = ga_line_fit(x, y)
a, b, rmse = history[-1]
print(f"Evolved: y = {a:.2f}x + {b:.2f}  (RMSE: {rmse:.4f})")
print(f"True:    y = 3.00x + 8.00")
# Evolved: y = 2.92x + 8.03  (RMSE: 0.2722)
# True:    y = 3.00x + 8.00

The result: Starting from 100 random guesses, the GA converges to $a \approx 2.92$ , $b \approx 8.03$ — close to the true values. No gradients, no closed-form solution — just evolution.

Scaling Up: The Travelling Salesman

Now let's tackle a problem where gradients don't exist. Given 25 cities at random positions, find the shortest route that visits every city exactly once and returns to the start.

The representation changes everything. A line is defined by two real numbers — you can average them. A route is a permutation of cities — you can't average city 3 and city 7 to get city 5. This means we need different crossover and mutation operators.

import numpy as np

def ordered_crossover(p1, p2):
    """Take a slice from parent 1, fill gaps from parent 2's order."""
    n = len(p1)
    start, end = sorted(np.random.choice(n, 2, replace=False))
    child = -np.ones(n, dtype=int)
    child[start:end] = p1[start:end]
    fill = [g for g in p2 if g not in child[start:end]]
    child[child == -1] = fill
    return child

def swap_mutate(route, rate):
    """Randomly swap cities with given probability per position."""
    route = route.copy()
    for i in range(len(route)):
        if np.random.random() < rate:
            j = np.random.randint(len(route))
            route[i], route[j] = route[j], route[i]
    return route

def ga_tsp(cities, pop_size=100, n_elites=20, mutation_rate=0.01,
           generations=500):
    """Evolve shortest route through all cities."""
    n = len(cities)

    def total_distance(route):
        shifts = np.roll(route, -1)
        return np.sqrt(((cities[route] - cities[shifts]) ** 2).sum(axis=1)).sum()

    pop = [np.random.permutation(n) for _ in range(pop_size)]
    history = []

    for gen in range(generations):
        distances = np.array([total_distance(r) for r in pop])
        fitness = 1.0 / distances
        history.append((pop[np.argmin(distances)].copy(), distances.min()))

        # Selection: elites + fitness-proportionate
        ranked = np.argsort(-fitness)
        elites = [pop[i].copy() for i in ranked[:n_elites]]
        probs = fitness / fitness.sum()
        idx = np.random.choice(pop_size, pop_size - n_elites, p=probs)
        np.random.shuffle(idx)

        # Crossover + mutation
        children = []
        for i in range(0, len(idx) - 1, 2):
            for p, q in [(idx[i], idx[i+1]), (idx[i+1], idx[i])]:
                children.append(swap_mutate(
                    ordered_crossover(pop[p], pop[q]), mutation_rate))
        children = children[:pop_size - n_elites]
        pop = elites + children

    return history

# 25 random cities
np.random.seed(42)
cities = np.random.rand(25, 2) * 200

history = ga_tsp(cities)
best_route, best_dist = history[-1]
print(f"Initial distance: {history[0][1]:.0f}")
print(f"Evolved distance: {best_dist:.0f}")
# Initial distance: 2043
# Evolved distance: 978

The result: The route distance drops from 2043 to 978 — a 52% improvement. Watch the GIF — chaotic crossings disappear as evolution untangles the route, generation by generation.

What Just Happened?

Both problems follow the same evolutionary loop:

Create a random population
Evaluate fitness (how good is each individual?)
Select the fittest parents
Breed offspring via crossover
Mutate for diversity
Repeat

The crucial difference is the representation — how we encode a solution as a "chromosome." This determines which operators make sense.

	Line Fitting	TSP
Individual	$(a, b)$ pair	City permutation
Gene	Real number	City index
Example	`[3.2, 7.9]`	`[0, 14, 3, 21, ...]`
Crossover	Average parents	Ordered crossover
Mutation	Gaussian noise	Swap two cities

Fitness: Measuring Quality

Fitness must be higher for better solutions. Both problems use inverse loss:

A line with RMSE 0.1 (fitness 10) is ten times "fitter" than one with RMSE 1.0 (fitness 1). This ratio drives the selection pressure.

Selection: Survival of the Fittest

We combine two mechanisms:

# Keep the top 20 unchanged (elitism)
elites = pop[ranked[:n_elites]]

# Fill remaining slots by fitness-proportionate sampling
probs = fitness / fitness.sum()
parents = pop[np.random.choice(pop_size, pop_size - n_elites, p=probs)]

Elitism guarantees we never lose our best solution. Roulette wheel selection gives fitter individuals a higher probability of being chosen as parents, but still allows less-fit individuals a chance — maintaining diversity.

Crossover: Where Representation Matters

This is the key insight. For line fitting, crossover is trivial — average the parents' genes:

child = (parent1 + parent2) / 2  # [3.5, 7.2] + [2.8, 8.1] → [3.15, 7.65]

For TSP, you can't average permutations — city 3.5 doesn't exist. Instead, we use ordered crossover (OX), which preserves relative ordering from both parents:

Copy a random slice from Parent 1
Fill remaining positions with cities from Parent 2, in order

Parent 1: [0, 1, 2, 3, 4, 5, 6, 7]
Parent 2: [2, 5, 0, 7, 1, 3, 6, 4]
                  ┌─────┐
Slice (pos 3-6):  [_, _, _, 3, 4, 5, _, _]

Fill from P2 (skip 3, 4, 5): [2, 0, 7, 1, 6]

Child:            [2, 0, 7, 3, 4, 5, 1, 6]  ← valid permutation

The child inherits a contiguous segment from Parent 1 and the relative ordering of remaining cities from Parent 2.

Mutation: Random Exploration

Without mutation, the population would eventually become clones of the fittest individual — stuck in a local optimum. Mutation injects diversity:

Line fitting: Gaussian noise — child *= 1 + rate × N(0, 1)
TSP: Swap two random cities — route[i], route[j] = route[j], route[i]

The mutation rate controls the exploration-exploitation trade-off. Too high and you're doing random search. Too low and you converge prematurely.

Going Deeper

Hyperparameters

Parameter	What it controls	Line fitting	TSP
`pop_size`	Search breadth	100	100
`n_elites`	Exploitation strength	20 (20%)	20 (20%)
`mutation_rate`	Exploration strength	0.1 (10%)	0.01 (1%)
`generations`	Computation budget	50	500

Notice the mutation rate difference: line fitting uses 10% because Gaussian noise is gentle (multiplicative perturbation), while TSP uses 1% because each swap can dramatically change the route.

GA vs Gradient Descent

For the line-fitting problem, we can compare GA with gradient descent directly:

def gd_line_fit(x, y, lr=0.1, iterations=50):
    """Gradient descent to fit y = ax + b."""
    a, b = 0.0, 0.0
    history = []
    for _ in range(iterations):
        preds = a * x + b
        error = preds - y
        a -= lr * 2 * (error * x).mean()
        b -= lr * 2 * error.mean()
        history.append((a, b, np.sqrt((error ** 2).mean())))
    return history

GD converges in ~15 iterations. GA takes ~30 generations with a population of 100 — that's 3,000 fitness evaluations vs 15. When gradients exist, use them. GA shines when they don't: combinatorial problems (TSP), black-box functions, or multi-modal landscapes where gradient descent gets stuck in local optima.

When NOT to Use Genetic Algorithms

Differentiable objectives — gradient descent is orders of magnitude faster
Small search spaces — just enumerate all solutions
Need optimality guarantees — GA only guarantees "good" solutions, not the best
Tight compute budget — evaluating the entire population each generation is expensive

Alternatives:

Gradient descent for differentiable problems
Simulated annealing for combinatorial optimisation with a single solution
Bayesian optimisation for expensive black-box functions
Ant colony optimisation for routing problems specifically

Deep Dive: The Papers

Holland (1975)

John Holland introduced genetic algorithms in Adaptation in Natural and Artificial Systems, published by the University of Michigan Press. Holland was interested in a fundamental question: how do complex adaptive systems — biological organisms, economies, ecosystems — improve over time?

His key insight was that adaptation can be studied as an algorithm. By abstracting the essentials of natural selection into a computational framework, he showed that a simple loop of selection, crossover, and mutation could solve optimisation problems across domains.

The Schema Theorem

Holland's most famous theoretical result is the Schema Theorem, sometimes called the Fundamental Theorem of Genetic Algorithms.

A schema is a template describing a subset of individuals. In a binary GA, the schema 1**0*1 matches any individual with a 1 in positions 0 and 5 and a 0 in position 3 (where * is a wildcard).

The theorem states:

Short, low-order, above-average schemata receive exponentially increasing trials in subsequent generations.

Formally:

Where:

$m(H, t)$ = instances of schema H at generation t
$f(H) / \bar{f}$ = schema fitness relative to population average
$\delta(H)$ = defining length (distance between outermost fixed positions)
$o(H)$ = order (number of fixed positions)
$p_c, p_m$ = crossover and mutation probabilities

In plain English: patterns that are short (small $\delta$ ), simple (small $o$ ), and above average ( $f(H) > \bar{f}$ ) spread exponentially through the population. This is why crossover works — it recombines short, successful building blocks.

Whitley (1994)

Darrell Whitley's A Genetic Algorithm Tutorial in Statistics and Computing provides an accessible theoretical treatment:

"A genetic algorithm is a model of machine learning which derives its behavior from a metaphor of the processes of evolution in nature. This is done by the creation within a machine of a population of individuals represented by chromosomes."

Whitley explains why different representations need different operators — exactly the issue we encountered with line fitting vs TSP. For permutation problems, he discusses the importance of operators that preserve feasibility (like our ordered crossover).

Our Implementation vs the Framework

GA concept	Line fitting	TSP
Chromosome	$(a, b)$ pair	City permutation
Fitness function	$1 / \text{RMSE}$	$1 / \text{distance}$
Selection	Elitism + roulette wheel	Elitism + roulette wheel
Crossover	Arithmetic average	Ordered crossover (OX)
Mutation	Gaussian perturbation	Swap mutation

Historical Context

Turing (1950) — Proposed "genetical or evolutionary search" for machine learning
Rechenberg (1965) — Evolution strategies for engineering optimisation
Holland (1975) — Formal GA framework with the Schema Theorem
De Jong (1975) — First systematic empirical study of GA behaviour
Goldberg (1989) — Popularised GAs in engineering with practical examples
Koza (1992) — Extended GAs to genetic programming (evolving programs)

Interactive Tools

Overfitting Explorer — See how model complexity affects generalisation, the same tradeoff GAs navigate when fitting functions

Backpropagation Demystified — gradient-based optimisation (the alternative for differentiable problems)
The EM Algorithm — iterative optimisation with convergence guarantees
From K-Means to GMM — another iterative optimisation algorithm
Maximum Likelihood Estimation — parameter estimation (what GA replaces here)

Try It Yourself

The interactive notebook includes exercises:

Noise sensitivity — Add increasing noise to the line data. How does GA cope compared to MLE?
More cities — Try 50 or 100 cities in the TSP. How does convergence change?
Tournament selection — Replace roulette wheel with tournament selection (pick k random individuals, keep the best). Compare convergence speed.
GA for neural nets — Use GA instead of backprop to train the XOR network from the backpropagation post. How many generations does it take?

Understanding genetic algorithms opens the door to evolutionary computation — a family of methods including evolution strategies, genetic programming, and neuroevolution. The core insight — that selection and variation can optimise anything with a fitness function — is one of the most versatile ideas in computer science.

Frequently Asked Questions

What is a genetic algorithm and when should I use it?

A genetic algorithm (GA) is an optimisation method inspired by natural selection. It maintains a population of candidate solutions that evolve through selection, crossover, and mutation over generations. Use GAs when you have a complex optimisation problem with no useful gradient information, many local optima, or a discrete/combinatorial search space.

What is the difference between genetic algorithms and gradient descent?

Gradient descent follows the gradient of the objective function to find a minimum, requiring the function to be differentiable. GAs are gradient-free: they only need to evaluate a fitness function, so they work on non-differentiable, discrete, or black-box problems. GAs explore more broadly but converge more slowly than gradient descent for smooth problems.

How do I choose the population size and mutation rate?

Larger populations explore more of the search space but are slower per generation. Start with 50-200 individuals for most problems. Mutation rate controls exploration: too low and the population converges prematurely; too high and it becomes random search. Typical rates are 1-5% per gene. Tune both by observing convergence curves.

What is crossover and why does it help?

Crossover combines parts of two parent solutions to create offspring, exploiting the idea that good solutions may share useful building blocks. Single-point crossover splits parents at one random position; uniform crossover mixes genes independently. Crossover accelerates search by combining useful traits from different parts of the population.

Can genetic algorithms solve the travelling salesman problem optimally?

GAs can find good solutions to TSP but do not guarantee optimality. For small instances (under 20-30 cities), exact methods are feasible. For larger instances, GAs with specialised operators (ordered crossover, inversion mutation) find near-optimal solutions efficiently. For the best results, combine GAs with local search heuristics like 2-opt.

Backpropagation Demystified: Neural Nets from First Principles

Berkan Sesen — Thu, 02 Apr 2026 08:52:44 +0000

Every modern deep learning framework — PyTorch, TensorFlow, JAX — does one thing brilliantly: it computes gradients for you. Call loss.backward() and millions of parameters update simultaneously. But what's actually happening under the hood?

Backpropagation is just the chain rule applied systematically through a computational graph. By the end of this post, you'll build a neural network from scratch — no frameworks, no autograd — and understand exactly how every weight gets updated. We'll start with something even simpler: watching gradient descent fit a line in real time.

The algorithm was popularised by Rumelhart, Hinton & Williams (1986) in their landmark Nature paper, though the mathematical foundations trace back to Linnainmaa (1970) and Werbos (1974).

Gradient Descent in 30 Seconds

Before we build a neural network, let's see what gradient descent actually looks like. We'll fit a line $y = ax + b$ to noisy data:

Watch how the line wobbles into place. At each step, we compute two gradients — $\partial L / \partial a$ and $\partial L / \partial b$ — and nudge the parameters in the direction that reduces the error. That's gradient descent.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# Generate noisy data: y = 2x + 1
np.random.seed(42)
x_data = np.linspace(0, 2, 30)
y_data = 2 * x_data + 1 + np.random.normal(0, 0.3, 30)

# Start with random parameters
a, b = 0.0, 0.0
lr = 0.1
history = [(a, b)]

# Gradient descent: 12 steps (converges fast!)
for _ in range(12):
    y_pred = a * x_data + b
    error = y_pred - y_data
    grad_a = (2 / len(x_data)) * np.sum(error * x_data)
    grad_b = (2 / len(x_data)) * np.sum(error)
    a -= lr * grad_a
    b -= lr * grad_b
    history.append((a, b))

# Animate
fig, ax = plt.subplots(figsize=(8, 5))
ax.scatter(x_data, y_data, alpha=0.6, label='Data')
line, = ax.plot([], [], 'r-', linewidth=2)
title_text = ax.set_title('')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_xlim(-0.1, 2.1)
ax.set_ylim(-0.5, 6)

def update(frame):
    a_f, b_f = history[frame]
    line.set_data(x_data, a_f * x_data + b_f)
    loss = np.mean((a_f * x_data + b_f - y_data) ** 2)
    title_text.set_text(f'Step {frame}: a={a_f:.2f}, b={b_f:.2f}, MSE={loss:.3f}')
    return line, title_text

anim = FuncAnimation(fig, update, frames=len(history), interval=500, blit=True)
plt.show()

But with just two parameters, computing gradients is trivial. A neural network might have thousands of parameters arranged in layers. How do you efficiently compute the gradient of the loss with respect to every weight? That's what backpropagation solves.

The Full Network: Learning XOR

Let's build something a single line can't solve. The XOR function returns 1 when inputs differ and 0 when they're the same:

Input 1	Input 2	XOR Output
0	0	0
0	1	1
1	0	1
1	1	0

No straight line can separate the 1s from the 0s — this is why Minsky & Papert (1969) argued that perceptrons were fundamentally limited. The fix? Add a hidden layer.

Here's a complete neural network that learns XOR from scratch:

import numpy as np

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def sigmoid_derivative(output):
    """Derivative of sigmoid given its output (not its input)."""
    return output * (1 - output)

class NeuralNetwork:
    def __init__(self, layer_sizes, learning_rate=0.5):
        self.lr = learning_rate
        self.weights = []
        self.biases = []
        for i in range(len(layer_sizes) - 1):
            w = np.random.randn(layer_sizes[i], layer_sizes[i + 1]) * 0.5
            b = np.random.randn(1, layer_sizes[i + 1]) * 0.5
            self.weights.append(w)
            self.biases.append(b)

    def forward(self, X):
        """Forward pass: compute output and cache activations."""
        self.activations = [X]
        current = X
        for w, b in zip(self.weights, self.biases):
            current = sigmoid(current @ w + b)
            self.activations.append(current)
        return current

    def backward(self, y_true):
        """Backward pass: compute gradients and update weights."""
        m = y_true.shape[0]
        # Output layer error
        delta = (self.activations[-1] - y_true) * sigmoid_derivative(self.activations[-1])

        # Propagate backwards through layers
        for i in range(len(self.weights) - 1, -1, -1):
            grad_w = self.activations[i].T @ delta / m
            grad_b = np.sum(delta, axis=0, keepdims=True) / m
            if i > 0:
                delta = (delta @ self.weights[i].T) * sigmoid_derivative(self.activations[i])
            self.weights[i] -= self.lr * grad_w
            self.biases[i] -= self.lr * grad_b

    def train(self, X, y, epochs=10000):
        """Train and return loss history."""
        losses = []
        for epoch in range(epochs):
            output = self.forward(X)
            loss = np.mean((y - output) ** 2)
            losses.append(loss)
            self.backward(y)
        return losses

# XOR data
X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
y = np.array([[0], [1], [1], [0]])

# Train: 2 inputs -> 4 hidden -> 1 output
np.random.seed(42)
nn = NeuralNetwork([2, 4, 1], learning_rate=2.0)
losses = nn.train(X, y, epochs=10000)

# Results
print("Predictions after training:")
for inputs, target in zip(X, y):
    pred = nn.forward(inputs.reshape(1, -1))
    print(f"  {inputs} -> {pred[0, 0]:.4f} (target: {target[0]})")

Predictions after training:
  [0 0] -> 0.0178 (target: 0)
  [0 1] -> 0.9840 (target: 1)
  [1 0] -> 0.9828 (target: 1)
  [1 1] -> 0.0174 (target: 0)

The network learns XOR perfectly. Let's see the loss curve and decision boundary:

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

# Loss curve
ax1.plot(losses)
ax1.set_xlabel('Epoch')
ax1.set_ylabel('MSE Loss')
ax1.set_title('Training Loss')
ax1.set_yscale('log')

# Decision boundary
xx, yy = np.meshgrid(np.linspace(-0.5, 1.5, 200),
                      np.linspace(-0.5, 1.5, 200))
grid = np.c_[xx.ravel(), yy.ravel()]
Z = nn.forward(grid).reshape(xx.shape)

ax2.contourf(xx, yy, Z, levels=50, cmap='RdBu_r', alpha=0.8)
ax2.scatter(X[:, 0], X[:, 1], c=y.ravel(), cmap='RdBu_r',
            edgecolors='black', s=200, linewidths=2, zorder=5)
ax2.set_xlabel('Input 1')
ax2.set_ylabel('Input 2')
ax2.set_title('XOR Decision Boundary')

plt.tight_layout()
plt.show()

The result: The loss drops over ~2,000 epochs, and the network learns a non-linear decision boundary that correctly separates the XOR pattern. No single line could do this.

What Just Happened?

The magic is in two functions: forward and backward. Let's trace exactly what happens for one training step.

The Forward Pass

The forward pass computes the network's prediction by pushing data through each layer:

def forward(self, X):
    self.activations = [X]
    current = X
    for w, b in zip(self.weights, self.biases):
        current = sigmoid(current @ w + b)
        self.activations.append(current)
    return current

For input [0, 1] (which should produce output 1):

Layer 1: Multiply inputs by weights, add bias, apply sigmoid to get hidden activations $h$
Layer 2: Multiply hidden activations by weights, add bias, apply sigmoid to get output $\hat{y}$

Each layer computes $\sigma(Wx + b)$ where $\sigma$ is the sigmoid function. Think of it as a pipeline: raw inputs, then learned features, then prediction. We cache every intermediate activation because we'll need them during the backward pass.

The Loss

After the forward pass, we measure how wrong we are:

This is the mean squared error — the same loss we used for fitting the line. The question is: how do we adjust every weight to reduce this loss?

The Backward Pass: Chain Rule in Action

The backward pass is where backpropagation happens. Let's trace the gradient for a single output weight $w_{jk}$ connecting hidden neuron $j$ to output neuron $k$ :

Three terms, each intuitive:

Term	Meaning	Formula
$\partial L / \partial \hat{y}_k$	How much does the loss change per unit change in prediction?	$\hat{y}_k - y_k$
$\partial \hat{y}_k / \partial z_k$	How much does the sigmoid output change per unit change in its input?	$\hat{y}_k(1 - \hat{y}_k)$
$\partial z_k / \partial w_{jk}$	How much does the weighted sum change per unit change in this weight?	$h_j$ (the hidden activation)

Multiply them together and you get the gradient. The code does this in two lines:

# Combine first two terms into "delta"
delta = (self.activations[-1] - y_true) * sigmoid_derivative(self.activations[-1])

# Multiply by the third term to get the weight gradient
grad_w = self.activations[i].T @ delta / m

How Gradients Flow Backwards

For the hidden layer, the chain gets longer. We need to propagate the error signal backwards through the output weights:

delta = (delta @ self.weights[i].T) * sigmoid_derivative(self.activations[i])

This is the "back" in backpropagation. The error at the output layer flows backwards through the weights, getting split and scaled at each connection. Each hidden neuron learns what fraction of the total error it contributed.

Think of it like assigning blame. If a hidden neuron had a large weight connecting it to the output, it "caused" more of the error and gets a bigger update. If its activation was near 0 or 1 (sigmoid saturating), the gradient shrinks — this is the vanishing gradient problem.

Going Deeper

The General Algorithm

Backpropagation for an $L$ -layer network:

Forward pass: For each layer $l = 1, \ldots, L$ , compute $z^{(l)} = W^{(l)} a^{(l-1)} + b^{(l)}$ and $a^{(l)} = \sigma(z^{(l)})$
Compute output error: $\delta^{(L)} = \nabla_a L \odot \sigma'(z^{(L)})$
Propagate backwards: For $l = L-1, \ldots, 1$ : $\delta^{(l)} = (W^{(l+1)})^T \delta^{(l+1)} \odot \sigma'(z^{(l)})$
Update weights: $W^{(l)} \leftarrow W^{(l)} - \eta \cdot \delta^{(l)} (a^{(l-1)})^T$

The key insight: each layer only needs the delta from the layer above it. Gradients flow like water down a chain of waterfalls — each layer passes its "error signal" to the layer below.

Hyperparameters

Parameter	What it controls	Guidance
Learning rate	Step size for weight updates	Too high: diverge; too low: slow. Start at 0.01-0.1
Hidden layer size	Model capacity	More neurons: more expressive, but slower and risk of overfitting
Epochs	How many times to see the data	Watch the loss curve; stop when it plateaus
Activation function	Non-linearity type	Sigmoid for this tutorial; ReLU is standard in modern networks

Why Sigmoid? (And Why Not)

We used sigmoid because it's the original activation from the 1986 paper and it makes the derivatives clean. But sigmoid has a problem: when the output is near 0 or 1, the derivative $\sigma'(z) = \sigma(z)(1 - \sigma(z))$ approaches zero. Gradients vanish, and deep layers stop learning.

Modern networks use ReLU ( $\max(0, x)$ ) which has a constant gradient of 1 for positive inputs. This simple change is one reason why deep learning took off in the 2010s.

When NOT to Use Manual Backpropagation

Production models — Use PyTorch, TensorFlow, or JAX. Their autograd is faster, numerically stable, and handles arbitrary computation graphs
Complex architectures — Attention, batch normalisation, skip connections all have tricky gradients; let the framework handle them
GPU training — Frameworks are optimised for hardware acceleration

Manual backprop is for understanding — once you know how it works, you'll debug framework code far more effectively.

Connection to Maximum Likelihood

Training a neural network with MSE loss is equivalent to maximum likelihood estimation under a Gaussian noise model. If we assume $y = f(x; \theta) + \epsilon$ where $\epsilon \sim \mathcal{N}(0, \sigma^2)$ , then minimising MSE maximises the likelihood. Similarly, cross-entropy loss corresponds to MLE under a Bernoulli model. The EM algorithm uses this same principle — its M-step is a weighted MLE, and neural nets are doing unweighted MLE via gradient descent.

Deep Dive: The Papers

Rumelhart, Hinton & Williams (1986)

The paper that put backpropagation on the map was "Learning representations by back-propagating errors", published in Nature. David Rumelhart, Geoffrey Hinton, and Ronald Williams showed that multi-layer networks could learn useful internal representations — directly addressing the limitations that Minsky & Papert (1969) had identified for single-layer perceptrons.

Their key contribution wasn't the algorithm itself (which had been discovered independently several times), but demonstrating that it worked on meaningful problems. They showed that hidden units could learn to represent features that weren't present in the input or output.

We describe a new learning procedure, back-propagation, for networks of neurone-like units. The procedure repeatedly adjusts the weights of the connections in the network so as to minimize a measure of the difference between the actual output vector of the net and the desired output vector.

The paper presents the generalised delta rule, which is exactly what we implemented:

Where $\delta_j$ is the error signal at unit $j$ . For output units, $\delta_j = (y_j - t_j) \sigma'(z_j)$ . For hidden units, $\delta_j = \sigma'(z_j) \sum_k w_{jk} \delta_k$ — the error propagated back from all units in the next layer.

The Earlier History

While Rumelhart et al. (1986) is the most cited, backpropagation was discovered independently by several researchers:

Linnainmaa (1970) published the reverse mode of automatic differentiation in his master's thesis — the mathematical foundation underlying backpropagation
Werbos (1974) described applying the chain rule to neural networks in his Harvard PhD thesis, predating Rumelhart by over a decade
Parker (1985) independently rediscovered the algorithm at Stanford

The algorithm's impact was delayed partly because computing power wasn't yet sufficient for large networks, and partly because the AI community was focused on symbolic approaches.

Our Implementation vs the Paper

Paper concept	Our code
Generalised delta rule	`delta = error * sigmoid_derivative(activation)`
Error propagation	`delta = (delta @ weights.T) * sigmoid_derivative(activation)`
Weight update	`weights -= lr * activations.T @ delta`
Differentiable activation	`sigmoid(x) = 1 / (1 + exp(-x))`

Historical Context

Linnainmaa (1970) — Reverse mode automatic differentiation
Werbos (1974) — First application to neural networks
Rumelhart, Hinton & Williams (1986) — Demonstrated practical learning of internal representations
LeCun et al. (1989) — Applied backprop to convolutional networks (digit recognition)
Hochreiter & Schmidhuber (1997) — LSTM, addressing the vanishing gradient problem
Hinton et al. (2006) — Deep belief networks, reigniting interest in deep learning
Baydin et al. (2018) — Modern review of automatic differentiation

Interactive Tools

Classification Metrics Calculator — Evaluate your neural network's predictions with precision, recall, and F1
Overfitting Explorer — Visualise the bias-variance tradeoff that regularisation in neural nets addresses

Maximum Likelihood Estimation from Scratch — Neural net training is MLE under the hood
The EM Algorithm: An Intuitive Guide — Another iterative optimisation algorithm, but for latent variables
From MLE to Bayesian Inference — What happens when you add priors to neural nets

Try It Yourself

The interactive notebook includes exercises:

Trace a weight update — Pick one specific weight. Compute its gradient by hand for one training example and verify it matches the code
Learning rate experiments — Try rates of 0.01, 0.5, 2.0, 10.0. What happens at each extreme?
Add a layer — Change the architecture from [2, 4, 1] to [2, 4, 4, 1]. Does it learn XOR faster or slower?
ReLU activation — Replace sigmoid with ReLU. What changes in the training dynamics?
Batch vs online — Modify the code to update weights after each example (online learning) instead of all four (batch). Compare convergence

Understanding backpropagation is the foundation of all modern deep learning. Whether you're training a simple classifier or a billion-parameter language model, the same principle applies: compute the prediction, measure the error, and propagate gradients backwards through every layer. The chain rule does the rest.

Frequently Asked Questions

What is backpropagation in simple terms?

Backpropagation is a method for computing how much each weight in a neural network contributed to the overall prediction error. It applies the chain rule of calculus systematically, starting from the output layer and working backwards through each hidden layer. Once every weight has a gradient, the optimiser nudges each one in the direction that reduces the error.

Why can a single-layer perceptron not learn XOR?

XOR requires a non-linear decision boundary because the positive examples (0,1) and (1,0) are not separable from the negative examples (0,0) and (1,1) by any straight line. A single-layer perceptron can only learn linear boundaries. Adding a hidden layer with a non-linear activation function gives the network the capacity to combine linear boundaries into non-linear regions.

What is the vanishing gradient problem?

When using sigmoid or tanh activations, the derivative approaches zero for very large or very small inputs. During backpropagation, these near-zero gradients are multiplied together across layers, causing the gradient signal to shrink exponentially as it flows backward. Deep layers receive almost no learning signal, effectively stalling training. ReLU activations largely solve this by maintaining a constant gradient of 1 for positive inputs.

Why do modern frameworks use autograd instead of manual backpropagation?

Manual backpropagation requires deriving and implementing gradients for every operation, which is error-prone and impractical for complex architectures. Autograd systems in PyTorch, TensorFlow, and JAX build a computational graph automatically and compute exact gradients through reverse-mode automatic differentiation. They also handle numerical stability, GPU acceleration, and arbitrary computation graphs that would be extremely difficult to implement by hand.

How do I choose the right learning rate?

The learning rate controls the step size of each weight update. Too high and the model diverges or oscillates; too low and training stalls. A common starting point is 0.01 for simple networks and 0.001 for deeper architectures. In practice, learning rate schedulers that reduce the rate over time or adaptive optimisers like Adam that adjust the rate per parameter are standard in modern deep learning.

MCMC Island Hopping: An Intuitive Guide to the Metropolis-Hastings Algorithm

Berkan Sesen — Wed, 01 Apr 2026 08:11:06 +0000

Imagine you're a politician touring a chain of islands. Each island has a different population, and you want to spend time on each island in proportion to its population — more time on crowded islands, less on quiet ones. The catch: you have no map. You can only see the island you're on and the one next door.

This is the core idea behind the Metropolis-Hastings algorithm, one of the most important algorithms in computational statistics. By the end of this post, you'll understand how a simple random walk can recover any target distribution, and you'll implement it from scratch.

The algorithm was originally developed for simulating equations of state in physics by Metropolis et al. (1953), and later generalised by Hastings (1970).

The Problem

You have a target distribution — a list of values representing how much time you should spend at each location. In our island example, the populations are:

Island	0	1	2	3	4	5	6
Population	2	3	1	5	8	2	9

You want to generate samples from this distribution. That is, if you record which island you visit at each time step, the histogram of visits should look like the population distribution — island 6 (population 9) should appear about 9 times more often than island 2 (population 1).

Why not just sample directly? In real applications, the distribution is often a complex posterior in thousands of dimensions. You can evaluate $P(x)$ at any point, but you can't normalise it or sample from it directly. MCMC sidesteps this entirely.

Quick Win: Run the Algorithm

Let's see the island hopper in action. Click the badge to open the interactive notebook:

Here's the complete implementation:

import numpy as np
import matplotlib.pyplot as plt

def metropolis_island_hopping(target, n_samples, start=0, burn_in=1000):
    """
    Metropolis-Hastings sampler for a discrete distribution using
    nearest-neighbour proposals (the 'island hopping' variant).

    Args:
        target: List of unnormalised target values (e.g. island populations)
        n_samples: Total number of steps (including burn-in)
        start: Starting island index
        burn_in: Steps to discard before collecting samples

    Returns:
        Array of visit counts for each island
    """
    n_islands = len(target)
    visit_counts = np.zeros(n_islands)
    current = start

    for step in range(n_samples):
        # PROPOSE: pick a neighbour uniformly at random
        if current == 0:
            proposed = 1
        elif current == n_islands - 1:
            proposed = n_islands - 2
        else:
            proposed = current + np.random.choice([-1, 1])

        # CORRECTION: account for asymmetric proposals at boundaries
        # q(proposed -> current) / q(current -> proposed)
        n_neighbours_current = 2 if 0 < current < n_islands - 1 else 1
        n_neighbours_proposed = 2 if 0 < proposed < n_islands - 1 else 1
        proposal_ratio = n_neighbours_current / n_neighbours_proposed

        # ACCEPT/REJECT: Metropolis-Hastings ratio
        acceptance_ratio = (target[proposed] / target[current]) * proposal_ratio

        if acceptance_ratio >= 1 or np.random.random() < acceptance_ratio:
            current = proposed

        # Record visit (after burn-in)
        if step >= burn_in:
            visit_counts[current] += 1

    return visit_counts

# Target distribution: island populations
target = [2, 3, 1, 5, 8, 2, 9]

# Run the sampler
visits = metropolis_island_hopping(target, n_samples=50_000, burn_in=1000)

The result: With few samples the histogram is noisy, but as the chain runs longer the orange bars converge to match the blue target — the sampler spends time on each island in proportion to its population, even though it only ever looks at one neighbour at a time.

What Just Happened?

The algorithm uses a beautifully simple loop: propose, evaluate, accept or reject. Let's walk through each part.

The Proposal: Pick a Neighbour

if current == 0:
    proposed = 1
elif current == n_islands - 1:
    proposed = n_islands - 2
else:
    proposed = current + np.random.choice([-1, 1])

Think of this as flipping a coin to decide whether to look left or right. At the boundaries (island 0 or island 6), there's only one direction to go. This asymmetry matters — we'll correct for it next.

The Proposal Correction

n_neighbours_current = 2 if 0 < current < n_islands - 1 else 1
n_neighbours_proposed = 2 if 0 < proposed < n_islands - 1 else 1
proposal_ratio = n_neighbours_current / n_neighbours_proposed

If you're on island 0 (one neighbour), you'll always propose island 1. But from island 1 (two neighbours), you only propose island 0 half the time. This asymmetry would bias the samples if we didn't correct for it.

The correction factor $q(\text{current} \mid \text{proposed}) / q(\text{proposed} \mid \text{current})$ ensures detailed balance — the fundamental property that makes MCMC work.

The Accept/Reject Decision

acceptance_ratio = (target[proposed] / target[current]) * proposal_ratio

if acceptance_ratio >= 1 or np.random.random() < acceptance_ratio:
    current = proposed

This is the heart of the algorithm. If the proposed island has a higher population, always move there. If it has a lower population, move with probability equal to the population ratio.

For example, if you're on island 4 (population 8) and propose island 3 (population 5):

You'd move with 62.5% probability. This means you sometimes visit less popular islands — which is essential for exploring the full distribution.

Why Burn-In?

if step >= burn_in:
    visit_counts[current] += 1

The starting position is arbitrary. The first few hundred steps are influenced by where you started, not by the target distribution. Discarding these "burn-in" samples removes this initial bias.

Going Deeper

The Acceptance Ratio

The general Metropolis-Hastings acceptance probability is:

Where:

$\pi(x')$ is the target distribution at the proposed state
$\pi(x)$ is the target distribution at the current state
$q(x' \mid x)$ is the probability of proposing $x'$ from $x$
$q(x \mid x')$ is the probability of proposing $x$ from $x'$

The ratio $\pi(x') / \pi(x)$ only needs unnormalised values — the normalising constant cancels out. This is why MCMC is so powerful: you never need to compute the full normalisation integral.

When the proposal distribution is symmetric ( $q(x' \mid x) = q(x \mid x')$ ), the correction factor disappears and we get the simpler Metropolis algorithm.

Hyperparameters

Parameter	What it controls	Guidance
`n_samples`	Total chain length	More is better; 10,000-100,000 typical
`burn_in`	Samples discarded from the start	10-20% of `n_samples`; check trace plots
`start`	Initial position	Shouldn't matter after burn-in, but starting near high-probability regions helps
Proposal width	How far each proposal jumps	Too small = slow exploration; too large = high rejection rate

Detailed Balance: Why It Works

The algorithm satisfies detailed balance:

This says: the probability of being at $x$ and transitioning to $x'$ equals the probability of being at $x'$ and transitioning to $x$ . Any Markov chain satisfying this condition has $\pi$ as its stationary distribution — so running the chain long enough guarantees convergence to the target.

When NOT to Use Metropolis-Hastings

High dimensions — Random walk proposals become inefficient; use Hamiltonian Monte Carlo (HMC) instead
Multimodal targets — The chain can get stuck in one mode; consider parallel tempering
Strong correlations — Axis-aligned proposals struggle; use adaptive proposals or Gibbs sampling
When you need speed — Variational inference is much faster (but approximate)

Alternatives:

Gibbs sampling for conditionally conjugate models
HMC/NUTS for continuous high-dimensional posteriors (what Stan uses)
Variational inference when speed matters more than exactness

Deep Dive: The Papers

Metropolis et al. (1953)

The original paper, "Equation of State Calculations by Fast Computing Machines", was written at Los Alamos National Laboratory. The authors — Nicholas Metropolis, Arianna and Marshall Rosenbluth, and Augusta and Edward Teller — needed to simulate the behaviour of molecules to compute thermodynamic properties.

Their key insight: instead of evaluating an intractable integral over all possible molecular configurations, generate representative configurations by random sampling and average over those.

The purpose of this paper is to describe a general method, suitable for fast computing machines, for investigating such properties as equations of state for substances consisting of interacting individual molecules.

The algorithm they proposed uses symmetric proposals only — a random perturbation of one particle's position. The acceptance rule was:

Where $\Delta E$ is the energy difference between the proposed and current configurations. In our island analogy, "energy" maps to the negative log of the population — lower energy means higher population.

Hastings (1970)

W.K. Hastings generalised the algorithm in "Monte Carlo Sampling Methods Using Markov Chains and Their Applications", allowing asymmetric proposal distributions. This is exactly the correction factor we implemented for the boundary islands.

Hastings showed that the acceptance probability:

guarantees detailed balance for any proposal distribution $q$ , not just symmetric ones. This seemingly small generalisation massively expanded the applicability of MCMC — it enabled samplers with directed proposals, mixture proposals, and eventually led to modern algorithms like HMC.

Our Implementation vs the Papers

Paper concept	Our code
Target $\pi(x)$	`target[i]` (island populations)
Proposal $q(x' \mid x)$	Uniform over neighbours (left/right)
Hastings correction	`proposal_ratio` for boundary asymmetry
Energy $\Delta E$	`-log(target[proposed]) + log(target[current])` (implicit)

Historical Context

Metropolis et al. (1953) — Original algorithm for molecular simulation
Hastings (1970) — Generalisation to asymmetric proposals
Geman & Geman (1984) — Gibbs sampling (special case of MH)
Gelfand & Smith (1990) — Brought MCMC to mainstream statistics
Neal (2011) — Hamiltonian Monte Carlo for efficient high-dimensional sampling

Try It Yourself

The interactive notebook includes exercises:

Trace plots — Plot the island index over time to visualise the random walk
Burn-in sensitivity — How does changing burn_in affect the result?
Chain length — Run with 1,000 vs 100,000 samples and compare accuracy
Cyclic islands — Modify the code so island 0 and island 6 are neighbours (wrap around)
Continuous target — Replace the discrete islands with a Gaussian distribution and a continuous proposal

Understanding Metropolis-Hastings opens the door to Bayesian inference, probabilistic programming (Stan, PyMC), and modern sampling methods like HMC. The core insight, that a carefully designed random walk converges to any target distribution, is one of the most powerful ideas in computational statistics.

Interactive Tools

Markov Chain Simulator — Visualise Markov chain dynamics, transition matrices, and stationary distributions
Bayes' Theorem Calculator — Compute posterior probabilities that MCMC approximates at scale

From MLE to Bayesian Inference — The Bayesian framework that MCMC makes computationally tractable
Hierarchical Bayesian Regression with PyMC — MCMC sampling in action with PyMC
Bayesian Survival Analysis with PyMC — Custom likelihoods sampled via MCMC
MCMC Poisson Mixture: Earthquake Regimes — MCMC for real-world regime detection

Frequently Asked Questions

What is MCMC and why do I need it?

MCMC (Markov Chain Monte Carlo) is a family of algorithms for sampling from probability distributions that are difficult to sample from directly. You need it when you want to do Bayesian inference but the posterior distribution has no closed-form solution, which is the case for most real-world models.

How does the Metropolis-Hastings algorithm work?

Start at a random point, propose a new point by adding random noise, and accept or reject the proposal based on how much more probable the new point is. If it is more probable, always accept. If less probable, accept with a probability equal to the ratio. Over many iterations, the chain of accepted points forms samples from the target distribution.

How do I know if my MCMC chain has converged?

Check trace plots for good mixing (the chain explores the full range of plausible values without getting stuck). Run multiple chains from different starting points and check that they agree (R-hat close to 1.0). Discard early samples as burn-in. If the chain looks like white noise around a stable mean, it has likely converged.

What is the burn-in period and how long should it be?

Burn-in is the initial period where the chain is still moving from its arbitrary starting point towards the high-probability region. These samples are discarded because they do not represent the target distribution. The length depends on the problem; monitor the trace plot and discard until the chain appears stationary. Typically 10-50% of total samples.

What is the acceptance rate and what should it be?

The acceptance rate is the fraction of proposed moves that are accepted. Too high (above 50%) means steps are too small and the chain explores slowly. Too low (below 15%) means steps are too large and most are rejected. For a multivariate target, aim for about 23% acceptance rate; for a univariate target, about 44%.

Hidden Markov Models: When Clusters Have Memory

Berkan Sesen — Tue, 31 Mar 2026 06:00:36 +0000

In the previous post, we clustered data into groups with K-Means and GMMs. Both methods assume observations are independent — each data point is analysed in isolation, with no regard for what came before it.

But time series data doesn't work that way. Today's stock market regime depends on yesterday's. A patient's health state today depends on their state last week. Weather tomorrow depends on weather today. When your clusters have memory, you need Hidden Markov Models.

By the end of this post, you'll implement the Forward and Viterbi algorithms from scratch, detect stock market regimes with a Gaussian HMM, and understand why an HMM produces 5x fewer regime switches than K-Means on the same data.

Markov Chains: Adding Memory to Randomness

Before we add the "hidden" part, let's understand Markov chains. A Markov chain is a sequence of random states where the next state depends only on the current state — not the entire history. This is the Markov property.

We define a chain with a transition matrix $\mathbf{A}$ , where $A_{ij} = P(S_{t+1} = j \mid S_t = i)$ .

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Patch

states = ['Sunny', 'Cloudy', 'Rainy']
A = np.array([
    [0.7, 0.2, 0.1],  # From Sunny
    [0.3, 0.4, 0.3],  # From Cloudy
    [0.2, 0.3, 0.5],  # From Rainy
])
pi = np.array([0.5, 0.3, 0.2])  # Initial state probabilities

def simulate_markov_chain(A, pi, n_steps, seed=42):
    """Simulate a Markov chain."""
    rng = np.random.default_rng(seed)
    state_seq = [rng.choice(len(pi), p=pi)]
    for _ in range(n_steps - 1):
        state_seq.append(rng.choice(len(pi), p=A[state_seq[-1]]))
    return np.array(state_seq)

chain = simulate_markov_chain(A, pi, n_steps=50)

colours_weather = ['#f39c12', '#95a5a6', '#3498db']
fig, ax = plt.subplots(figsize=(12, 2.5))
for t, s in enumerate(chain):
    ax.bar(t, 1, color=colours_weather[s], width=1.0, edgecolor='white', linewidth=0.5)
ax.set_xlabel('Time step')
ax.set_yticks([])
ax.set_title('Simulated Weather Markov Chain')
ax.legend(handles=[Patch(color=c, label=s) for c, s in zip(colours_weather, states)],
          loc='upper right', ncol=3)
plt.tight_layout()
plt.show()

Notice the temporal clustering — rainy days tend to follow rainy days, sunny streaks persist. This structure is exactly what K-Means and GMMs cannot capture, because they assume observations are i.i.d.

Quick Win: The Forward and Viterbi Algorithms

Now add the "hidden" part. In an HMM, the Markov chain runs in a hidden layer — you can't observe the states directly. You only see emissions that depend on the hidden state.

The umbrella example: You can't see the weather (hidden), but each day you observe whether your colleague carries an umbrella (emitted).

Click the badge to open the interactive notebook:

An HMM has three sets of parameters $\lambda = (\mathbf{A}, \mathbf{B}, \boldsymbol{\pi})$ :

A — Transition matrix: $P(\text{next state} \mid \text{current state})$
B — Emission probabilities: $P(\text{observation} \mid \text{state})$
$\boldsymbol{\pi}$ — Initial state distribution

# Hidden states: Sunny (0), Rainy (1)
# Observations: No umbrella (0), Umbrella (1)

A_hmm = np.array([
    [0.7, 0.3],  # Sunny → Sunny, Sunny → Rainy
    [0.4, 0.6],  # Rainy → Sunny, Rainy → Rainy
])

B_hmm = np.array([
    [0.9, 0.1],  # Sunny: P(no umbrella), P(umbrella)
    [0.2, 0.8],  # Rainy: P(no umbrella), P(umbrella)
])

pi_hmm = np.array([0.6, 0.4])  # P(Sunny), P(Rainy) at t=0

Given these parameters and the observation sequence [Umbrella, Umbrella, No umbrella, Umbrella, ...], we want to answer two questions:

How likely is this sequence? (Forward algorithm)
What was the most likely weather each day? (Viterbi algorithm)

The Forward Algorithm

The forward variable $\alpha_t(i)$ is the joint probability of seeing the first $t$ observations and being in state $i$ at time $t$ :

def forward_algorithm(observations, A, B, pi):
    """
    Forward algorithm for HMMs.

    Returns:
        alpha: (T, K) forward variables
        log_likelihood: log P(O | λ)
    """
    T = len(observations)
    K = len(pi)
    alpha = np.zeros((T, K))

    # Initialisation: α_1(i) = π_i * B_i(O_1)
    alpha[0] = pi * B[:, observations[0]]

    # Recursion: α_t(j) = [Σ_i α_{t-1}(i) * A_ij] * B_j(O_t)
    for t in range(1, T):
        alpha[t] = (alpha[t-1] @ A) * B[:, observations[t]]

    # Termination: P(O|λ) = Σ_i α_T(i)
    log_likelihood = np.log(alpha[-1].sum())
    return alpha, log_likelihood

The Viterbi Algorithm

While the Forward algorithm gives probabilities, Viterbi finds the single most likely state sequence using dynamic programming — replacing the sum with a max:

def viterbi(observations, A, B, pi):
    """Viterbi algorithm — find the most likely state sequence."""
    T = len(observations)
    K = len(pi)

    log_A = np.log(A)
    log_B = np.log(B)
    log_pi = np.log(pi)

    delta = np.zeros((T, K))
    psi = np.zeros((T, K), dtype=int)  # backpointers

    # Initialisation
    delta[0] = log_pi + log_B[:, observations[0]]

    # Recursion: max instead of sum
    for t in range(1, T):
        for j in range(K):
            scores = delta[t-1] + log_A[:, j]
            psi[t, j] = np.argmax(scores)
            delta[t, j] = scores[psi[t, j]] + log_B[j, observations[t]]

    # Backtrack
    best_path = np.zeros(T, dtype=int)
    best_path[-1] = np.argmax(delta[-1])
    for t in range(T-2, -1, -1):
        best_path[t] = psi[t+1, best_path[t+1]]

    return best_path, delta[-1, best_path[-1]]

Let's run both on our umbrella observations:

observations = np.array([1, 1, 0, 1, 0, 0, 1, 1, 1, 0])

alpha, log_lik = forward_algorithm(observations, A_hmm, B_hmm, pi_hmm)
state_probs = alpha / alpha.sum(axis=1, keepdims=True)
best_path, _ = viterbi(observations, A_hmm, B_hmm, pi_hmm)

fig, axes = plt.subplots(2, 1, figsize=(10, 5), sharex=True)
days = range(1, len(observations) + 1)

ax = axes[0]
ax.bar(days, state_probs[:, 0], color='#f39c12', alpha=0.8, label='P(Sunny)')
ax.bar(days, state_probs[:, 1], bottom=state_probs[:, 0], color='#3498db', alpha=0.8, label='P(Rainy)')
ax.set_ylabel('Probability')
ax.set_title('Forward Algorithm: Filtered State Probabilities')
ax.legend(loc='upper right')

ax = axes[1]
viterbi_colours = ['#f39c12' if s == 0 else '#3498db' for s in best_path]
ax.bar(days, [1]*len(days), color=viterbi_colours, edgecolor='white', linewidth=0.5)
obs_labels = ['Umbrella' if o else 'No umbrella' for o in observations]
for t, (o, s) in enumerate(zip(obs_labels, best_path)):
    ax.text(t+1, 0.5, o.split()[0][:3], ha='center', va='center', fontsize=8, fontweight='bold')
ax.set_ylabel('State')
ax.set_xlabel('Day')
ax.set_title('Viterbi: Most Likely State Sequence')
ax.set_yticks([])

plt.tight_layout()
plt.show()

The result: On days with umbrellas, the model assigns high probability to rainy weather. The Viterbi path shows the most likely sequence — note how it respects temporal persistence (rainy days cluster together) rather than flipping on every observation.

What Just Happened?

Three Tasks with HMMs

HMMs support three fundamental tasks, first formalised by Rabiner (1989):

Task	Question	Algorithm	Complexity
Evaluation	How likely is this sequence?	Forward	$\mathcal{O}(TK^2)$
Decoding	What's the most likely state sequence?	Viterbi	$\mathcal{O}(TK^2)$
Learning	What parameters best explain the data?	Baum-Welch (EM)	$\mathcal{O}(TK^2)$ per iteration

We implemented Tasks 1 and 2. Task 3 — learning the parameters from data — uses the Baum-Welch algorithm, which is EM applied to HMMs. If you understood EM for coin tosses and EM for GMMs, the Baum-Welch algorithm follows the same pattern:

E-step: Compute the expected state occupancies (forward-backward algorithm)
M-step: Re-estimate A, B, and $\boldsymbol{\pi}$ using these expected counts

From Discrete to Continuous: Gaussian HMMs

Our umbrella example had discrete emissions (umbrella or not). For continuous data like stock returns, we replace the emission matrix B with Gaussian emission densities — each hidden state $k$ emits observations from $\mathcal{N}(\boldsymbol{\mu}_k, \boldsymbol{\Sigma}_k)$ .

This is where HMMs meet GMMs: a Gaussian HMM is essentially a GMM where the mixture weights change over time according to the transition matrix.

Stock Market Regime Detection

Now let's apply a Gaussian HMM to a real problem: detecting market regimes (bull, bear, sideways) in S&P 500 data.

import yfinance as yf
from sklearn.preprocessing import StandardScaler
from hmmlearn.hmm import GaussianHMM

ticker = yf.Ticker('^GSPC')
df = ticker.history(start='2018-01-01', end='2024-01-01')

df['Return'] = df['Close'].pct_change()
df = df[df['Volume'] > 0]
df['Log_Volume'] = np.log(df['Volume'])
df = df.dropna()
df = df[np.isfinite(df['Return']) & np.isfinite(df['Log_Volume'])]

features = df[['Return', 'Log_Volume']].values
scaler = StandardScaler()
X_scaled = scaler.fit_transform(features)

# Fit a 3-state Gaussian HMM
model = GaussianHMM(n_components=3, covariance_type='full', n_iter=200, random_state=42)
model.fit(X_scaled)

hidden_states = model.predict(X_scaled)    # Viterbi decoding
posteriors = model.predict_proba(X_scaled)  # Forward-backward posteriors

# Visualise regimes
regime_returns = [features[hidden_states == k, 0].mean() for k in range(3)]
order = np.argsort(regime_returns)
regime_names = {order[0]: 'Bear', order[1]: 'Sideways', order[2]: 'Bull'}
regime_colours = {order[0]: '#e74c3c', order[1]: '#f39c12', order[2]: '#2ecc71'}

fig, axes = plt.subplots(2, 1, figsize=(12, 7), sharex=True,
                         gridspec_kw={'height_ratios': [3, 1]})

ax = axes[0]
dates = df.index
close = df['Close'].values
for k in range(3):
    mask = hidden_states == k
    ax.scatter(dates[mask], close[mask], c=regime_colours[k], s=3, alpha=0.7, label=regime_names[k])
ax.set_ylabel('S&P 500 Close')
ax.set_title('S&P 500 Market Regimes Detected by HMM')
ax.legend(loc='upper left', markerscale=5)
ax.grid(True, alpha=0.3)

ax = axes[1]
bottom = np.zeros(len(dates))
for k in [order[2], order[1], order[0]]:
    ax.fill_between(dates, bottom, bottom + posteriors[:, k],
                    color=regime_colours[k], alpha=0.8)
    bottom += posteriors[:, k]
ax.set_ylabel('P(regime)')
ax.set_xlabel('Date')
ax.set_title('Posterior Regime Probabilities')
ax.set_ylim(0, 1)

plt.tight_layout()
plt.show()

The HMM identifies clear regimes: the COVID crash of March 2020 shows up as a sharp bear regime, while the sustained rallies of 2019 and 2021 are classified as bull. The posterior probabilities (bottom panel) show how confident the model is about each regime assignment.

K-Means vs HMM: The Key Comparison

Now for the punchline. Let's run K-Means on the same features (daily return + log volume) and compare:

from sklearn.cluster import KMeans

km = KMeans(n_clusters=3, random_state=42, n_init=10)
km_labels = km.fit_predict(X_scaled)

fig, axes = plt.subplots(2, 1, figsize=(12, 7), sharex=True)

# K-Means
ax = axes[0]
for k in range(3):
    mask = km_labels == k
    ax.scatter(dates[mask], close[mask], c=km_colours[k], s=3, alpha=0.7, label=km_names[k])
ax.set_title('K-Means: Clusters Based on Features Only (ignores time)')
ax.legend(loc='upper left', markerscale=5)
ax.grid(True, alpha=0.3)

# HMM
ax = axes[1]
for k in range(3):
    mask = hidden_states == k
    ax.scatter(dates[mask], close[mask], c=regime_colours[k], s=3, alpha=0.7, label=regime_names[k])
ax.set_title('HMM: Regimes with Temporal Dependency')
ax.legend(loc='upper left', markerscale=5)
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

The difference is striking:

K-Means produces scattered, noisy assignments with 538 regime switches — it classifies each day independently based on that day's return and volume
HMM produces smooth, persistent regimes with only 105 switches — it knows that a bull market today makes a bull market tomorrow more likely

This is the power of the transition matrix. K-Means treats every observation as independent. The HMM's transition matrix encodes the prior that regimes persist — which is exactly how markets behave.

Going Deeper

The Unified View: K-Means → GMM → HMM

These three algorithms form a natural progression, each relaxing one assumption:

	K-Means	GMM	HMM
Responsibilities	Hard (0 or 1)	Soft (probabilities)	Soft (probabilities)
Data assumption	i.i.d.	i.i.d.	Sequentially dependent
Cluster shape	Spherical (isotropic)	Elliptical (full covariance)	Elliptical (full covariance)
Objective	Distortion $J$	Log-likelihood $\ell$	Log-likelihood $\ell$
Fitting	Lloyd's (EM, hard)	EM (soft)	Baum-Welch (EM for sequences)

K-Means → GMM: relax hard to soft assignments.
GMM → HMM: add sequential dependency between observations.

All three use variations of the EM algorithm. The M-step is always weighted MLE — the difference is how the weights (responsibilities) are computed in the E-step.

Hyperparameters

Parameter	What it controls	Guidance
`n_components`	Number of hidden states	2-5 typical for financial data; use BIC to select
`covariance_type`	Emission covariance structure	`full` for correlated features, `diag` for independent
`n_iter`	Maximum EM iterations	200-1000; check convergence
Feature selection	What the model "sees"	Returns, volume, volatility — domain knowledge matters

When NOT to Use HMMs

No temporal structure — If your observations genuinely are independent, an HMM adds unnecessary complexity. Use a GMM.
Long-range dependencies — The Markov property means the model only looks one step back. For patterns spanning weeks or months, consider RNNs or Transformers.
Many hidden states — The $\mathcal{O}(TK^2)$ complexity makes large state spaces expensive. For K > 10, consider variational methods.
Non-stationary dynamics — HMMs assume the transition matrix is fixed over time. Regime-switching models or online learning may be more appropriate for evolving markets.

Deep Dive: The Papers

Rabiner (1989)

The definitive HMM tutorial is Rabiner, L.R. (1989) "A Tutorial on Hidden Markov Models and Selected Applications in Speech Recognition," Proceedings of the IEEE, 77(2), 257-286.

Originally written for the speech recognition community, this paper has become the standard reference across all fields that use HMMs. Rabiner's clear presentation of the three fundamental problems — evaluation, decoding, and learning — and their efficient solutions made HMMs accessible to a generation of researchers.

"Given three or more barrels of brandy and a set of wine glasses, one of the simplest hidden Markov models for the process of tasting the brandy is..."
— Rabiner (1989), using a brandy-tasting analogy to introduce HMMs

Baum and Welch (1970)

The EM algorithm for HMMs was developed by Baum, L.E. et al. (1970) "A Maximization Technique Occurring in the Statistical Analysis of Probabilistic Functions of Markov Chains," Annals of Mathematical Statistics, 41(1), 164-171.

This preceded Dempster, Laird and Rubin's general EM framework (1977) by seven years. The Baum-Welch algorithm is a specific instance of EM where the E-step uses the forward-backward algorithm to compute expected state occupancies and transition counts.

Our Implementation vs the Papers

Paper concept	Our code
Forward variable $\alpha_t(i)$	`alpha[t] = (alpha[t-1] @ A) * B[:, observations[t]]`
Viterbi variable $\delta_t(j)$	`delta[t, j] = scores[psi[t, j]] + log_B[j, observations[t]]`
Transition matrix A	`model.transmat_` (learned by hmmlearn)
Emission densities B	`model.means_`, `model.covars_` (Gaussian)
Three tasks	Forward (eval), Viterbi (decode), Baum-Welch (learn)

Historical Context

Baum & Welch (1970) — EM for HMMs (Baum-Welch algorithm)
Viterbi (1967) — Efficient decoding via dynamic programming
Dempster, Laird & Rubin (1977) — General EM framework (unifies Baum-Welch with other methods)
Rabiner (1989) — The definitive HMM tutorial
Bishop's PRML Chapter 13 — Modern textbook treatment

Try It Yourself

The interactive notebook includes exercises:

Number of states — Fit HMMs with 2-5 states and use BIC to choose the best K
Different stocks — Apply the HMM to a volatile stock (Tesla) vs a stable one (Johnson & Johnson)
Implement Baum-Welch — Extend the forward algorithm with a backward pass and implement the parameter updates from scratch
Regime-conditional trading — Compute Sharpe ratios within each regime and design a regime-aware strategy
Discrete HMM — Discretise returns into bins and compare a discrete HMM with the Gaussian version

Interactive Tools

Markov Chain Simulator — Visualise Markov chain dynamics, transition matrices, and stationary distributions in the browser

From K-Means to GMM: Hard vs Soft Clustering — The i.i.d. clustering methods that HMMs extend with temporal dependency
Gaussian Mixture Models: EM in Practice — GMMs are the i.i.d. special case of Gaussian HMMs
The EM Algorithm: An Intuitive Guide — The Baum-Welch algorithm is EM applied to HMMs
Maximum Likelihood Estimation from Scratch — The M-step in Baum-Welch is weighted MLE, just like in GMMs

Frequently Asked Questions

What is a Hidden Markov Model and when should I use it?

An HMM models sequential data where the observed outputs depend on hidden (latent) states that evolve over time according to a Markov chain. Use HMMs when your data has temporal structure and you believe it is generated by switching between a small number of regimes, such as speech phonemes, market regimes, or biological sequence states.

What is the difference between an HMM and a standard Markov chain?

In a Markov chain, the states are directly observed. In an HMM, the states are hidden and you only observe noisy emissions from each state. The challenge is to infer the hidden state sequence from the observations, which requires specialised algorithms like the forward-backward algorithm and the Viterbi algorithm.

What are the three fundamental HMM problems?

Evaluation (forward algorithm): what is the probability of an observed sequence given the model? Decoding (Viterbi algorithm): what is the most likely hidden state sequence? Learning (Baum-Welch/EM): what model parameters best explain the observed data? These three problems cover the main use cases of HMMs.

How do I choose the number of hidden states?

Use information criteria (BIC/AIC) to compare models with different numbers of states, or use domain knowledge if you know how many regimes exist. Too few states underfit (missing real structure); too many overfit (splitting genuine states). Cross-validation on held-out sequences can also guide the choice.

What are the limitations of HMMs?

HMMs assume the Markov property (the future depends only on the current state, not the history), which is too restrictive for some applications. They also assume discrete hidden states and parametric emission distributions. For richer models, consider recurrent neural networks, conditional random fields, or hierarchical HMMs.

From K-Means to GMM: Hard vs Soft Clustering

Berkan Sesen — Mon, 30 Mar 2026 07:54:19 +0000

You have a pile of unlabelled data and you want to find groups in it. K-Means is the algorithm everyone reaches for first — it's fast, simple, and usually works. But it makes a bold assumption: every data point belongs to exactly one cluster. No uncertainty, no hedging.

What happens to a point sitting right between two clusters? K-Means forces a choice. Gaussian Mixture Models (GMMs) offer an alternative — soft assignments that express how uncertain we are. By the end of this post, you'll implement K-Means from scratch, see why it's secretly a special case of the EM algorithm, and understand exactly when soft clustering beats hard clustering.

Quick Win: K-Means from Scratch

Let's cluster some data. Click the badge to open the interactive notebook:

Here's K-Means in 30 lines — Lloyd's algorithm:

import numpy as np

def kmeans(X, K, max_iter=100, seed=42):
    """
    K-Means clustering (Lloyd's algorithm).

    Args:
        X: (N, D) array of data points
        K: number of clusters
        max_iter: maximum iterations
        seed: random seed for reproducibility

    Returns:
        centroids: (K, D) final cluster centres
        labels: (N,) hard cluster assignments
        history: list of (centroids, labels) at each iteration
        distortions: list of objective function values
    """
    rng = np.random.default_rng(seed)
    N, D = X.shape

    # Step 1: Initialise centroids by picking K random data points
    indices = rng.choice(N, K, replace=False)
    centroids = X[indices].copy()

    history = []
    distortions = []

    for iteration in range(max_iter):
        # Step 2: ASSIGN — each point to nearest centroid
        distances = np.linalg.norm(X[:, None] - centroids[None, :], axis=2)
        labels = np.argmin(distances, axis=1)

        # Compute distortion (objective function)
        distortion = sum(np.sum((X[labels == k] - centroids[k])**2) for k in range(K))
        distortions.append(distortion)
        history.append((centroids.copy(), labels.copy()))

        # Step 3: UPDATE — move centroids to cluster means
        new_centroids = np.array([X[labels == k].mean(axis=0) for k in range(K)])

        # Check convergence
        if np.allclose(centroids, new_centroids):
            break

        centroids = new_centroids

    return centroids, labels, history, distortions

Run it on three synthetic clusters:

import matplotlib.pyplot as plt

rng = np.random.default_rng(42)
n_per_cluster = 100
cluster_1 = rng.multivariate_normal([2, 2], [[0.8, 0.4], [0.4, 0.3]], n_per_cluster)
cluster_2 = rng.multivariate_normal([7, 7], [[0.5, -0.3], [-0.3, 0.8]], n_per_cluster)
cluster_3 = rng.multivariate_normal([2, 8], [[0.3, 0.0], [0.0, 0.6]], n_per_cluster)
X = np.vstack([cluster_1, cluster_2, cluster_3])

centroids, labels, history, distortions = kmeans(X, K=3)

colours = ['#e74c3c', '#3498db', '#2ecc71']

fig, axes = plt.subplots(1, 2, figsize=(12, 5))

ax = axes[0]
for k in range(3):
    mask = labels == k
    ax.scatter(X[mask, 0], X[mask, 1], c=colours[k], alpha=0.5, s=20, label=f'Cluster {k+1}')
ax.scatter(centroids[:, 0], centroids[:, 1], c='black', marker='X', s=200, zorder=5, label='Centroids')
ax.set_xlabel('$x_1$'); ax.set_ylabel('$x_2$')
ax.set_title('K-Means Result (K=3)')
ax.legend(); ax.grid(True, alpha=0.3)

ax = axes[1]
ax.plot(range(1, len(distortions) + 1), distortions, 'k-o', markersize=6)
ax.set_xlabel('Iteration'); ax.set_ylabel('Distortion $J$')
ax.set_title('K-Means Objective Decreases at Each Step')
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

The result: K-Means finds three clean clusters in about 5 iterations. The distortion (sum of squared distances to centroids) drops sharply and converges. Each point is assigned to exactly one cluster — no ambiguity.

What Just Happened?

K-Means repeats two steps until nothing changes:

Step 1: Assign Points to Nearest Centroid

distances = np.linalg.norm(X[:, None] - centroids[None, :], axis=2)
labels = np.argmin(distances, axis=1)

For each data point, compute the Euclidean distance to every centroid and pick the closest one. This creates hard responsibilities — each point gets a label of 0, 1, or 2, with no in-between.

Step 2: Move Centroids to Cluster Means

new_centroids = np.array([X[labels == k].mean(axis=0) for k in range(K)])

Each centroid moves to the average position of all points assigned to it. Think of it as a "centre of gravity" — the centroid is pulled toward its members.

Why It Converges

Each iteration can only decrease (or maintain) the distortion. Assigning a point to a closer centroid reduces the sum of squared distances. Moving a centroid to the cluster mean minimises the within-cluster variance. Since distortion is bounded below by zero and decreases monotonically, the algorithm must converge.

Let's watch this happen step by step:

The centroids start at random data points and quickly settle into the cluster centres. By iteration 3-4, the assignments are essentially fixed.

from matplotlib.animation import FuncAnimation, PillowWriter

fig, ax = plt.subplots(figsize=(8, 6))

def update(frame):
    ax.clear()
    c, l = history[frame]
    for k in range(3):
        mask = l == k
        ax.scatter(X[mask, 0], X[mask, 1], c=colours[k], alpha=0.4, s=15)
    ax.scatter(c[:, 0], c[:, 1], c='black', marker='X', s=200, zorder=5)
    ax.set_title(f'K-Means: Iteration {frame + 1}')
    ax.set_xlim(X[:, 0].min()-1, X[:, 0].max()+1)
    ax.set_ylim(X[:, 1].min()-1, X[:, 1].max()+1)
    ax.grid(True, alpha=0.3)

anim = FuncAnimation(fig, update, frames=len(history), interval=600)
plt.show()

The Boundary Problem

K-Means makes a hard decision for every point. But what about a point sitting exactly between two centroids? It gets forced into one cluster with 100% confidence, even though a 50/50 split would be more honest.

# Distance ratios reveal ambiguous points
distances = np.linalg.norm(X[:, None] - centroids[None, :], axis=2)
sorted_dists = np.sort(distances, axis=1)
ambiguity = 1 - (sorted_dists[:, 0] / sorted_dists[:, 1])

fig, axes = plt.subplots(1, 2, figsize=(12, 5))

ax = axes[0]
for k in range(3):
    mask = labels == k
    ax.scatter(X[mask, 0], X[mask, 1], c=colours[k], alpha=0.5, s=20)
ax.scatter(centroids[:, 0], centroids[:, 1], c='black', marker='X', s=200, zorder=5)
ax.set_title('K-Means: Hard Assignments')
ax.set_xlabel('$x_1$'); ax.set_ylabel('$x_2$'); ax.grid(True, alpha=0.3)

ax = axes[1]
sc = ax.scatter(X[:, 0], X[:, 1], c=ambiguity, cmap='RdYlGn_r', s=20, alpha=0.7)
ax.scatter(centroids[:, 0], centroids[:, 1], c='black', marker='X', s=200, zorder=5)
plt.colorbar(sc, ax=ax, label='Ambiguity')
ax.set_title('How Uncertain Is Each Assignment?')
ax.set_xlabel('$x_1$'); ax.set_ylabel('$x_2$'); ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

The red points near the cluster boundaries are assigned with high confidence by K-Means, but they're actually ambiguous. This is the fundamental limitation of hard clustering.

Going Deeper

The Distortion Objective

K-Means minimises the distortion (also called inertia or within-cluster sum of squares):

Where $r_{nk} \in \{0, 1\}$ is the hard responsibility — 1 if point $n$ is assigned to cluster $k$ , 0 otherwise. The centroid $\boldsymbol{\mu}_k$ is the mean of all points assigned to cluster $k$ .

K-Means Is EM with Hard Assignments

Here's the key insight: K-Means is a special case of the EM algorithm where the responsibilities are binary (0 or 1) instead of continuous probabilities.

To see this, write K-Means in EM notation:

def kmeans_as_em(X, K, max_iter=100, seed=42):
    """K-Means expressed as EM with hard (0/1) responsibilities."""
    rng = np.random.default_rng(seed)
    N, D = X.shape

    indices = rng.choice(N, K, replace=False)
    centroids = X[indices].copy()

    for iteration in range(max_iter):
        # E-STEP: Compute hard responsibilities r(n,k) ∈ {0, 1}
        distances = np.linalg.norm(X[:, None] - centroids[None, :], axis=2)
        r = np.zeros((N, K))
        r[np.arange(N), np.argmin(distances, axis=1)] = 1.0

        # M-STEP: Update centroids using responsibilities
        Nk = r.sum(axis=0)
        new_centroids = (r.T @ X) / Nk[:, None]

        if np.allclose(centroids, new_centroids):
            break
        centroids = new_centroids

    return centroids, r

Compare this with the GMM E-step where $r_{nk}$ is a probability between 0 and 1. The only difference is whether the responsibility is hard or soft.

As Bishop (2006) puts it: K-Means corresponds to the zero-temperature limit of the GMM. If you shrink all covariance matrices toward zero, the Gaussian components become infinitely peaked, and the soft responsibilities snap to 0 or 1.

K-Means vs GMM: Side by Side

Let's fit both to the same data and compare:

from matplotlib.patches import Ellipse

def draw_ellipse(ax, mean, cov, colour, n_std=2):
    vals, vecs = np.linalg.eigh(cov)
    angle = np.degrees(np.arctan2(vecs[1, 1], vecs[0, 1]))
    width, height = 2 * n_std * np.sqrt(vals)
    ell = Ellipse(xy=mean, width=width, height=height, angle=angle,
                  edgecolor=colour, facecolor='none', linewidth=2, linestyle='--')
    ax.add_patch(ell)

fig, axes = plt.subplots(1, 2, figsize=(13, 5.5))

# K-Means
ax = axes[0]
for k in range(3):
    mask = labels == k
    ax.scatter(X[mask, 0], X[mask, 1], c=colours[k], alpha=0.5, s=20)
ax.scatter(centroids[:, 0], centroids[:, 1], c='black', marker='X', s=200, zorder=5)
ax.set_title('K-Means: Hard Assignments\n(each point → exactly one cluster)')
ax.set_xlabel('$x_1$'); ax.set_ylabel('$x_2$'); ax.grid(True, alpha=0.3)

# GMM
ax = axes[1]
gmm_labels = np.argmax(gmm_resp, axis=1)
for k in range(3):
    mask = gmm_labels == k
    alphas = gmm_resp[mask, k]
    ax.scatter(X[mask, 0], X[mask, 1], c=colours[k], alpha=0.3 + 0.5*alphas, s=20)
    draw_ellipse(ax, gmm_means[k], gmm_covs[k], colours[k])
ax.scatter(gmm_means[:, 0], gmm_means[:, 1], c='black', marker='X', s=200, zorder=5)
ax.set_title('GMM: Soft Assignments\n(each point → probability per cluster)')
ax.set_xlabel('$x_1$'); ax.set_ylabel('$x_2$'); ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

Notice two key differences:

Cluster shapes — K-Means produces spherical (Voronoi) boundaries. The GMM captures the elongated shape of each cluster with covariance ellipses.
Transparency — In the GMM plot, points near boundaries are more transparent, reflecting their split responsibilities.

The Assumptions Progression

K-Means, GMMs, and Hidden Markov Models form a natural progression. Each relaxes one assumption from its predecessor:

	K-Means	GMM	HMM
Responsibilities	Hard (0 or 1)	Soft (probabilities)	Soft (probabilities)
Data assumption	i.i.d.	i.i.d.	Sequentially dependent
Cluster shape	Spherical (isotropic)	Elliptical (full covariance)	Elliptical (full covariance)
Objective	Distortion $J$	Log-likelihood $\ell$	Log-likelihood $\ell$
Fitting	Lloyd's algorithm	EM	Baum-Welch (EM for sequences)

K-Means → GMM: relax hard to soft assignments. GMM → HMM: add sequential dependency between observations.

Choosing K: The Elbow Method

K-Means requires you to pick $K$ upfront. The elbow method plots the final distortion against $K$ and looks for the "bend":

K_range = range(1, 8)
final_distortions = []

for K in K_range:
    _, _, _, dists = kmeans(X, K=K)
    final_distortions.append(dists[-1])

plt.figure(figsize=(7, 4))
plt.plot(list(K_range), final_distortions, 'ko-', markersize=8)
plt.axvline(x=3, color='red', linestyle='--', alpha=0.5, label='True K=3')
plt.xlabel('Number of Clusters (K)')
plt.ylabel('Final Distortion $J$')
plt.title('Elbow Method for Choosing K')
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()

Distortion always decreases with more clusters (at $K = N$ it's zero). The "elbow" is where adding another cluster gives diminishing returns. For our data, the bend at $K = 3$ is clear.

For a more principled approach, GMMs offer the Bayesian Information Criterion (BIC), which penalises model complexity — see the GMM tutorial for details.

Hyperparameters

Parameter	What it controls	Guidance
`K`	Number of clusters	Use elbow method, silhouette score, or BIC
`max_iter`	Maximum iterations	100-300 usually sufficient
`seed` / initialisation	Starting centroids	K-Means++ (default in scikit-learn) is more robust than random
`n_init`	Number of random restarts	10 is typical; keeps the best result

When NOT to Use K-Means

Non-spherical clusters — K-Means assumes isotropic distance. Elongated or curved clusters get split incorrectly. Use GMMs or DBSCAN.
Unequal cluster sizes — K-Means tends to produce roughly equal-sized clusters, even when the true clusters differ dramatically in size.
You need uncertainty — If a point's cluster membership matters (e.g., for downstream decisions), use a GMM to get probabilities.
Sequential data — If your observations have a natural ordering (time series), the i.i.d. assumption breaks. Use HMMs instead.

When K-Means wins: It's fast ( $\mathcal{O}(NKD)$ per iteration), scales to millions of points, and when clusters are roughly spherical and well-separated, it gives the same answer as the GMM with far less computation.

Deep Dive: The Papers

Lloyd (1982) and the History of K-Means

The algorithm we call "K-Means" was first described by Stuart Lloyd in a 1957 Bell Labs technical report, but wasn't formally published until 1982 as "Least Squares Quantization in PCM" in IEEE Transactions on Information Theory. Lloyd was working on pulse-code modulation — converting analogue signals to digital — and needed to find optimal quantisation levels that minimised reconstruction error.

Meanwhile, James MacQueen independently described and named the algorithm "K-Means" in a 1967 conference paper, giving it the name we still use today.

The Distortion Objective

Lloyd's algorithm minimises:

Each iteration alternates between:

Minimise $J$ over $r_{nk}$ (with $\boldsymbol{\mu}_k$ fixed) — this is the assignment step
Minimise $J$ over $\boldsymbol{\mu}_k$ (with $r_{nk}$ fixed) — this is the update step

Setting the derivative of $J$ with respect to $\boldsymbol{\mu}_k$ to zero gives:

which is just the mean of the assigned points — exactly what our code computes.

Bishop's Unified View (PRML Chapter 9)

Bishop's Pattern Recognition and Machine Learning (2006), Chapter 9, provides the clearest treatment of the K-Means–GMM connection. He shows that K-Means arises as a limiting case of the GMM when we:

Fix all covariance matrices to $\epsilon \mathbf{I}$ (spherical, equal variance)
Take the limit $\epsilon \to 0$

As $\epsilon$ shrinks, the Gaussian responsibilities $r_{nk} = \pi_k \mathcal{N}(\mathbf{x}_n \mid \boldsymbol{\mu}_k, \epsilon \mathbf{I})$ become increasingly peaked, converging to 0 or 1 — exactly the hard assignments of K-Means.

"We can also derive the K-means algorithm as a particular limit of the EM algorithm for Gaussian mixtures."
— Bishop, PRML (2006), §9.3.2

Our Implementation vs the Papers

Paper concept	Our code
Distortion $J$	`distortion = sum(np.sum((X[labels == k] - centroids[k])**2) ...)`
Hard responsibilities $r_{nk}$	`labels = np.argmin(distances, axis=1)`
Centroid update $\boldsymbol{\mu}_k$	`X[labels == k].mean(axis=0)`
Lloyd's assign-update loop	The `for iteration` loop

Try It Yourself

The interactive notebook includes exercises:

Initialisation sensitivity — Run K-Means with 10 different random seeds. How often does it find the "right" clusters? Try K-Means++ initialisation.
Non-spherical clusters — Create elongated clusters. How does K-Means handle them compared to GMM?
Old Faithful — Apply both K-Means and GMM to the Old Faithful geyser data (2D: duration + waiting time). Compare the cluster boundaries.
Soft K-Means — Modify the code to use soft assignments: $r_{nk} = e^{-\beta d_{nk}^2} / \sum_j e^{-\beta d_{nj}^2}$ . What happens as $\beta \to \infty$ ?
Compare with scikit-learn — Verify your implementation matches sklearn.cluster.KMeans on the same data.

Interactive Tools

Distribution Explorer — Visualise the Gaussian and other distributions that underpin mixture models
Classification Metrics Calculator — Evaluate clustering quality with metrics like precision and recall

The EM Algorithm: An Intuitive Guide — K-Means is EM with hard assignments; the coin-toss example shows why soft assignments work better
Gaussian Mixture Models: EM in Practice — The soft-clustering generalisation of K-Means, applied to Old Faithful geyser data
Maximum Likelihood Estimation from Scratch — The M-step in both K-Means and GMM is a form of (weighted) MLE

Frequently Asked Questions

What is the difference between hard and soft clustering?

Hard clustering assigns each data point to exactly one cluster (as in K-Means), while soft clustering assigns a probability of belonging to each cluster (as in GMMs). Soft clustering is more informative because it captures uncertainty: points near cluster boundaries get meaningful probabilities for multiple clusters rather than being forced into one.

When should I use GMM instead of K-Means?

Use GMM when your clusters have different shapes, sizes, or orientations, since K-Means assumes spherical clusters of similar size. GMM is also better when you need membership probabilities rather than hard assignments, or when clusters overlap significantly. K-Means is faster and simpler, so use it when clusters are roughly spherical and well-separated.

How do I choose the number of components in a GMM?

Use information criteria like BIC or AIC, which balance fit quality against model complexity. Fit GMMs with different numbers of components and pick the one with the lowest BIC. You can also use silhouette scores or domain knowledge. Unlike K-Means, GMM can also use the likelihood to compare models directly.

Can GMM handle clusters of different shapes?

Yes. Each component in a GMM has its own covariance matrix, which can model elliptical clusters at any orientation. By contrast, K-Means implicitly assumes spherical clusters. You can also constrain the covariance matrices (diagonal, spherical, tied) to reduce the number of parameters when data is limited.

Why does K-Means sometimes give poor results?

K-Means fails when clusters are non-spherical, have very different sizes, or overlap heavily. It is also sensitive to initialisation (solved partially by K-Means++) and to outliers, which can pull centroids away from the true cluster centres. If K-Means gives poor results, GMM or density-based methods like DBSCAN may be more appropriate.

From MLE to Bayesian Inference: Why Your Estimate Needs a Prior

Berkan Sesen — Sun, 29 Mar 2026 08:24:37 +0000

In the MLE tutorial, we estimated a coin's bias by finding the single parameter value that maximises the likelihood. Flip a coin 3 times, get 3 heads, and MLE says $\hat{\theta} = 1.0$ — the coin always lands heads. That feels wrong. With only 3 flips, we shouldn't be certain of anything.

The problem isn't the likelihood — it's that MLE gives you a point estimate with no way to express doubt. Bayesian inference fixes this by computing an entire distribution over parameter values, weighted by how plausible each value is given both the data and your prior knowledge. By the end of this post, you'll implement Bayesian updating from scratch, understand conjugate priors, and see why a 99% accurate medical test can still be wrong 98% of the time.

Quick Win: A Coin Flip with a Prior

Let's revisit the coin flip from the MLE post, but this time we'll incorporate a prior belief. Suppose you think the coin is probably fair, but you're not certain.

import numpy as np
import matplotlib.pyplot as plt
from scipy import stats

# Prior: Beta(2, 2) — mild belief the coin is roughly fair
alpha_prior, beta_prior = 2, 2

# Data: 7 heads out of 10 flips
z, N = 7, 10

# Posterior: Beta(alpha + z, beta + N - z)
alpha_post = alpha_prior + z
beta_post = beta_prior + (N - z)

# Plot prior, likelihood, and posterior
theta = np.linspace(0, 1, 500)
prior = stats.beta.pdf(theta, alpha_prior, beta_prior)
likelihood = theta**z * (1 - theta)**(N - z)
likelihood /= likelihood.max()  # normalise for plotting
posterior = stats.beta.pdf(theta, alpha_post, beta_post)

fig, axes = plt.subplots(3, 1, figsize=(8, 8), sharex=True)

axes[0].fill_between(theta, prior, alpha=0.3, color='blue')
axes[0].plot(theta, prior, 'b-', linewidth=2)
axes[0].set_ylabel('Density')
axes[0].set_title('Prior: Beta(2, 2)')

axes[1].fill_between(theta, likelihood, alpha=0.3, color='green')
axes[1].plot(theta, likelihood, 'g-', linewidth=2)
axes[1].set_ylabel('Likelihood (normalised)')
axes[1].set_title(f'Likelihood: {z} heads in {N} flips')

axes[2].fill_between(theta, posterior, alpha=0.3, color='red')
axes[2].plot(theta, posterior, 'r-', linewidth=2)
axes[2].axvline(z/N, color='green', linestyle='--', label=f'MLE: {z/N:.2f}')
axes[2].axvline(alpha_post/(alpha_post + beta_post), color='red', linestyle='--',
                label=f'Posterior mean: {alpha_post/(alpha_post + beta_post):.2f}')
axes[2].set_xlabel('θ (coin bias)')
axes[2].set_ylabel('Density')
axes[2].set_title(f'Posterior: Beta({alpha_post}, {beta_post})')
axes[2].legend()

plt.tight_layout()
plt.show()

print(f"MLE:            θ = {z/N:.3f}")
print(f"Posterior mean:  θ = {alpha_post/(alpha_post + beta_post):.3f}")
print(f"95% credible interval: [{stats.beta.ppf(0.025, alpha_post, beta_post):.3f}, "
      f"{stats.beta.ppf(0.975, alpha_post, beta_post):.3f}]")

The three panels show: the prior (your initial belief), the likelihood (what the data says), and the posterior (your updated belief). The posterior sits between the prior and the likelihood — a compromise weighted by how much data you have.

Notice two things the MLE can't give you:

The posterior mean (0.643) is pulled toward 0.5 compared to the MLE (0.700), because the prior nudges us toward fairness.
The 95% credible interval tells you exactly where the true bias probably lies — no bootstrapping or asymptotic arguments needed.

What Just Happened?

Bayes' Rule: The One Equation

Everything in Bayesian inference flows from a single equation:

Each term has a name:

Term	Name	What it means
$p(\theta \mid D)$	Posterior	What we believe about $\theta$ after seeing data
$p(D \mid \theta)$	Likelihood	How probable is this data for a given $\theta$
$p(\theta)$	Prior	What we believed about $\theta$ before seeing data
$p(D)$	Evidence	The total probability of the data across all $\theta$ values

As I explained in a CrossValidated answer years ago:

MLE treats $p(\theta)/p(D)$ as a constant and seeks a single point $\hat{\theta}$ that maximises the likelihood. Bayesian estimation fully calculates the posterior $p(\theta \mid D)$ , treating $\theta$ as a random variable. We put in probability density functions and get out probability density functions — not a single point.

The Evidence (Marginal Likelihood)

The denominator $p(D)$ is the trickiest part. It ensures the posterior integrates to 1:

This integral sums the likelihood over all possible parameter values, weighted by the prior. For most models, this integral has no closed-form solution — which is exactly why methods like MCMC sampling exist.

But for some lucky combinations of prior and likelihood, the integral works out analytically. These are called conjugate priors.

The Disease Diagnosis Surprise

Before diving into conjugate priors, here's why priors matter so much — even with excellent data. This example, from Kruschke (2015), is one of the most counterintuitive results in all of probability.

Suppose a disease affects 1 in 1,000 people. A diagnostic test has:

99% hit rate: if you have the disease, the test is positive 99% of the time
5% false alarm rate: if you're healthy, the test is still positive 5% of the time

You test positive. What's the probability you actually have the disease?

# Prior: base rate of the disease
p_disease = 0.001
p_healthy = 1 - p_disease

# Likelihood: test characteristics
p_positive_given_disease = 0.99   # hit rate
p_positive_given_healthy = 0.05   # false alarm rate

# Evidence: total probability of testing positive
p_positive = (p_positive_given_disease * p_disease +
              p_positive_given_healthy * p_healthy)

# Posterior: Bayes' Rule
p_disease_given_positive = (p_positive_given_disease * p_disease) / p_positive

print(f"Prior probability of disease:  {p_disease:.4f} ({p_disease*100:.1f}%)")
print(f"Posterior after positive test:  {p_disease_given_positive:.4f} "
      f"({p_disease_given_positive*100:.1f}%)")

The answer: only 1.9%. Despite the 99% hit rate, you almost certainly don't have the disease. Why? Because the disease is so rare (prior = 0.001) that the vast majority of positive results come from false alarms among healthy people.

This is Bayes' rule in action: the posterior depends on both the likelihood (test accuracy) and the prior (base rate). Ignoring the prior — as MLE effectively does — would lead to wildly overconfident conclusions.

Conjugate Priors: The Beta-Binomial

Why Conjugate Priors Exist

The evidence integral $p(D) = \int p(D \mid \theta) \, p(\theta) \, d\theta$ is usually intractable. But if we choose a prior from the right mathematical family, the posterior has a closed-form solution. Such a prior is called conjugate to the likelihood.

For the Bernoulli/Binomial likelihood, the conjugate prior is the Beta distribution.

The Beta Distribution

The Beta distribution lives on $[0, 1]$ — perfect for modelling a probability parameter:

where $B(\alpha, \beta)$ is the Beta function (a normalising constant). The two parameters control the shape:

fig, axes = plt.subplots(2, 3, figsize=(12, 6))

params = [(1, 1), (2, 2), (5, 5), (0.5, 0.5), (2, 5), (20, 20)]
titles = ['Uniform\nα=1, β=1', 'Mild\nα=2, β=2', 'Moderate\nα=5, β=5',
          'Jeffreys\nα=0.5, β=0.5', 'Asymmetric\nα=2, β=5', 'Strong\nα=20, β=20']

theta = np.linspace(0.001, 0.999, 500)
for ax, (a, b), title in zip(axes.flat, params, titles):
    ax.plot(theta, stats.beta.pdf(theta, a, b), 'b-', linewidth=2)
    ax.fill_between(theta, stats.beta.pdf(theta, a, b), alpha=0.2)
    ax.set_title(title)
    ax.set_xlabel('θ')
    ax.set_xlim(0, 1)

plt.tight_layout()
plt.show()

Think of $\alpha$ and $\beta$ as pseudo-counts: $\alpha - 1$ imaginary heads and $\beta - 1$ imaginary tails from a previous experiment. A flat $\text{Beta}(1, 1)$ means you've seen zero imaginary data — total ignorance.

The Conjugate Update

Here's the key result. If our prior is $\text{Beta}(\alpha, \beta)$ and we observe $z$ heads in $N$ flips, the posterior is:

Combining the exponents:

This is another Beta distribution:

That's the entire update rule. The prior had $\alpha - 1$ pseudo-heads and $\beta - 1$ pseudo-tails. The data contributed $z$ real heads and $N - z$ real tails. The posterior simply adds them together. We never needed to compute the evidence integral — the conjugacy handled it.

Summary Statistics

From the posterior $\text{Beta}(\alpha', \beta')$ where $\alpha' = \alpha + z$ and $\beta' = \beta + N - z$ :

def bayesian_coin_summary(alpha_prior, beta_prior, z, N):
    """Posterior summary for the Beta-Binomial model."""
    alpha_post = alpha_prior + z
    beta_post = beta_prior + (N - z)

    posterior_mean = alpha_post / (alpha_post + beta_post)
    posterior_mode = (alpha_post - 1) / (alpha_post + beta_post - 2)
    mle = z / N
    ci_low = stats.beta.ppf(0.025, alpha_post, beta_post)
    ci_high = stats.beta.ppf(0.975, alpha_post, beta_post)

    print(f"Prior:           Beta({alpha_prior}, {beta_prior})")
    print(f"Data:            {z} heads in {N} flips")
    print(f"Posterior:       Beta({alpha_post}, {beta_post})")
    print(f"MLE:             {mle:.4f}")
    print(f"MAP estimate:    {posterior_mode:.4f}")
    print(f"Posterior mean:  {posterior_mean:.4f}")
    print(f"95% CI:          [{ci_low:.4f}, {ci_high:.4f}]")

# The coin from the MLE post: 73 heads in 100 flips
bayesian_coin_summary(2, 2, 73, 100)

With 100 data points, the prior barely matters — the posterior mean is nearly identical to the MLE. The prior's influence fades as data accumulates.

The Tug-of-War: Prior vs. Data

More Data, Less Prior Influence

The posterior is always a compromise between the prior and the likelihood. With more data, the likelihood wins:

fig, axes = plt.subplots(1, 3, figsize=(14, 4))
theta = np.linspace(0, 1, 500)
alpha_prior, beta_prior = 10, 10  # Prior centred at 0.5

for ax, (z, N) in zip(axes, [(1, 4), (5, 20), (25, 100)]):
    alpha_post = alpha_prior + z
    beta_post = beta_prior + (N - z)

    prior = stats.beta.pdf(theta, alpha_prior, beta_prior)
    posterior = stats.beta.pdf(theta, alpha_post, beta_post)
    likelihood = theta**z * (1 - theta)**(N - z)
    likelihood = likelihood / likelihood.max() * posterior.max()

    ax.plot(theta, prior, 'b--', linewidth=1.5, label='Prior', alpha=0.7)
    ax.plot(theta, likelihood, 'g:', linewidth=1.5, label='Likelihood', alpha=0.7)
    ax.fill_between(theta, posterior, alpha=0.3, color='red')
    ax.plot(theta, posterior, 'r-', linewidth=2, label='Posterior')
    ax.axvline(z/N, color='green', linestyle='--', alpha=0.4)
    ax.set_title(f'{z} heads in {N} flips (25%)')
    ax.set_xlabel('θ')
    ax.legend(fontsize=8)

plt.tight_layout()
plt.show()

With 4 flips, the prior pulls the posterior toward 0.5. With 100 flips, the posterior clusters tightly around the MLE at 0.25. As Kruschke puts it: "The compromise favours the prior to the extent that the prior distribution is sharply peaked and the data are few."

Stronger Prior, More Resistance

A sharper prior requires more data to overcome:

fig, axes = plt.subplots(1, 3, figsize=(14, 4))
theta = np.linspace(0, 1, 500)
z, N = 3, 10

priors = [(1, 1, 'Flat (α=1, β=1)'),
          (5, 5, 'Moderate (α=5, β=5)'),
          (30, 30, 'Strong (α=30, β=30)')]

for ax, (a, b, title) in zip(axes, priors):
    alpha_post = a + z
    beta_post = b + (N - z)

    prior = stats.beta.pdf(theta, a, b)
    posterior = stats.beta.pdf(theta, alpha_post, beta_post)

    ax.plot(theta, prior, 'b--', linewidth=1.5, label='Prior', alpha=0.7)
    ax.fill_between(theta, posterior, alpha=0.3, color='red')
    ax.plot(theta, posterior, 'r-', linewidth=2, label='Posterior')
    ax.axvline(z/N, color='green', linestyle='--', alpha=0.5, label=f'MLE: {z/N:.1f}')
    mean = alpha_post / (alpha_post + beta_post)
    ax.axvline(mean, color='red', linestyle='--', alpha=0.5,
               label=f'Mean: {mean:.2f}')
    ax.set_title(title)
    ax.set_xlabel('θ')
    ax.legend(fontsize=8)

plt.tight_layout()
plt.show()

With a flat prior, the posterior mode equals the MLE (0.30). With a strong prior centred at 0.5, ten flips barely move the posterior. This is rational: a sharp prior represents genuine previous knowledge that we'd be reluctant to abandon without overwhelming evidence.

MLE Is Bayesian with a Flat Prior

When the prior is $\text{Beta}(1, 1)$ — the uniform distribution — the posterior mode (MAP estimate) simplifies to:

MLE is just Bayesian inference with a uniform prior. It's not "prior-free" — it implicitly assumes every parameter value is equally plausible before seeing data. When you use MLE, you are making a prior assumption, whether you realise it or not.

Sequential Updating

One powerful property of Bayesian inference: you can update beliefs incrementally as new data arrives. Today's posterior becomes tomorrow's prior.

Watch how the posterior (red) starts as a vague prior (blue) and sharpens with each batch of data, converging toward the MLE (green dashed line).

alpha, beta_param = 2, 2
theta = np.linspace(0, 1, 500)

batches = [(3, 5, 'Batch 1: 3/5 heads'),
           (6, 10, 'Batch 2: 6/10 heads'),
           (15, 20, 'Batch 3: 15/20 heads')]

fig, axes = plt.subplots(1, len(batches) + 1, figsize=(16, 3.5))

axes[0].fill_between(theta, stats.beta.pdf(theta, alpha, beta_param),
                     alpha=0.3, color='blue')
axes[0].plot(theta, stats.beta.pdf(theta, alpha, beta_param), 'b-', linewidth=2)
axes[0].set_title(f'Prior\nBeta({alpha}, {beta_param})')
axes[0].set_xlabel('θ')

total_z, total_N = 0, 0
for i, (z_i, n_i, label) in enumerate(batches):
    alpha += z_i
    beta_param += (n_i - z_i)
    total_z += z_i
    total_N += n_i

    axes[i + 1].fill_between(theta, stats.beta.pdf(theta, alpha, beta_param),
                              alpha=0.3, color='red')
    axes[i + 1].plot(theta, stats.beta.pdf(theta, alpha, beta_param),
                     'r-', linewidth=2)
    axes[i + 1].axvline(total_z / total_N, color='green', linestyle='--', alpha=0.5)
    axes[i + 1].set_title(f'After {label}\nBeta({alpha}, {beta_param})')
    axes[i + 1].set_xlabel('θ')

plt.tight_layout()
plt.show()

The order doesn't matter: updating with Batch 1 then Batch 2 gives the same posterior as Batch 2 then Batch 1, as long as the data are independent. This is data-order invariance (Kruschke, 2015, Section 5.2.1).

When Bayesian Inference Gets Hard

The Beta-Binomial worked out neatly because of conjugacy. But most real models don't have conjugate priors. Consider Gaussian Mixture Models: you'd need to integrate over means, variances, and mixture weights simultaneously. The evidence integral becomes impossibly complex.

Two main approaches handle this:

The EM Algorithm — sidesteps the full posterior by alternating between estimating hidden variables and maximising the likelihood. It gives point estimates, not full posteriors, but it's computationally efficient. See the EM tutorial.
MCMC Sampling — draws samples from the posterior without computing the evidence integral. The Metropolis-Hastings algorithm generates a chain of samples that, after enough steps, faithfully represent the posterior distribution.

Deep Dive: The Foundations

Bayes' Original Paper (1763)

Thomas Bayes (1702-1761) was a Presbyterian minister and mathematician in England. His theorem was published posthumously in 1763, edited by his friend Richard Price. Bayes himself probably didn't fully grasp the ramifications — it was Pierre-Simon Laplace (1749-1827) who independently rediscovered and extensively developed Bayesian methods (Kruschke, 2015).

Kruschke's "Doing Bayesian Data Analysis"

Our treatment follows Chapter 5 of Kruschke (2015). His key pedagogical insight is presenting Bayes' rule as spatial attention in a joint probability table: conditioning on observed data means restricting attention to one row and renormalising. The disease diagnosis example, the coin bias estimation, and the prior-vs-data visualisations all come from this chapter.

Bishop's Pattern Recognition and Machine Learning

Bishop (2006), Chapter 2 provides the mathematical framework for the Beta-Binomial conjugate pair:

The Beta prior $\text{Beta}(\alpha, \beta)$ has mean $\alpha / (\alpha + \beta)$ and can be interpreted as $\alpha + \beta - 2$ prior observations
The sequential update property means the prior can encode genuine previous data, not just subjective belief
As the number of observations grows, the posterior concentrates around the MLE — Bayesian and frequentist approaches converge asymptotically

A Concise Comparison

For a concise summary of MLE vs. Bayesian estimation, see my CrossValidated answer. The key distinction: MLE treats $p(\theta)/p(D)$ as a constant and finds a single point $\hat{\theta}$ . Bayesian estimation fully calculates the posterior $p(\theta \mid D)$ , treating $\theta$ as a random variable. We put in distributions and get out distributions.

The trade-off is complexity. As noted in that answer, dealing with the evidence integral $p(D) = \int p(D \mid \theta) \, p(\theta) \, d\theta$ leads directly to the concept of conjugate priors: for a given likelihood, we choose a prior form that allows us to carry out the integration analytically.

Try It Yourself

The interactive notebook includes exercises:

The Sunrise Problem — You've seen the sun rise every day of your life. Use Bayesian inference to compute the probability it rises tomorrow. How does your prior change the answer?
Normal-Normal Conjugacy — The Normal distribution is conjugate to itself. With known variance, derive the posterior for the mean given a Normal prior
MAP vs. Posterior Mean — When do the MAP estimate and posterior mean disagree most? Create an example with a skewed Beta prior
Drug Trial — A new drug cures 8 out of 10 patients. Your prior from previous studies is $\text{Beta}(3, 7)$ (30% success rate). What does the posterior tell you?
Breaking Conjugacy — Replace the Beta prior with a mixture of two Betas and show the posterior is no longer Beta. Approximate it using grid methods

Interactive Tools

Bayes' Theorem Calculator — Run the disease diagnosis example and other Bayesian calculations in the browser
Distribution Explorer — Visualise Beta, Normal, and other distributions used in Bayesian updating

Maximum Likelihood Estimation from Scratch — The starting point: MLE gives point estimates by maximising the likelihood alone
The EM Algorithm — EM uses MLE internally but handles the case where some data is hidden
MCMC Island Hopping — When the posterior integral is intractable, sample from it instead

Frequently Asked Questions

What is the difference between MLE and Bayesian inference?

MLE finds the single parameter value that makes the observed data most probable, giving you a point estimate. Bayesian inference combines the data with prior beliefs to produce a full probability distribution over parameters (the posterior). The posterior tells you not just the best estimate but how uncertain you should be about it.

What is a prior and how do I choose one?

A prior represents your beliefs about the parameter before seeing data. If you have domain knowledge, encode it (for example, a coin is probably close to fair). If not, use weakly informative priors that rule out implausible values without strongly influencing the result. As you collect more data, the prior matters less and the posterior is dominated by the likelihood.

When does Bayesian inference give different results from MLE?

The difference is largest when you have little data, because the prior has more influence. With large datasets, the Bayesian posterior concentrates around the MLE and the two approaches give nearly identical point estimates. Bayesian inference is most valuable when data is scarce, when you need uncertainty quantification, or when you want to incorporate prior knowledge.

What is a conjugate prior?

A conjugate prior is one where the posterior has the same functional form as the prior, making the Bayesian update analytically tractable. For example, a Beta prior with a Binomial likelihood gives a Beta posterior. Conjugate priors are convenient for hand calculations, but modern computational methods (MCMC) allow you to use any prior you want.

Is Bayesian inference always better than MLE?

Not necessarily. MLE is simpler, faster, and sufficient for many applications, especially with large datasets. Bayesian inference adds value when uncertainty quantification matters, when data is limited, when you have meaningful prior information, or when you need to make decisions under uncertainty. The computational cost is also higher.

Gaussian Mixture Models: The EM Algorithm in Practice

Berkan Sesen — Sat, 28 Mar 2026 09:41:12 +0000

You have a dataset that clearly contains distinct groups, but nobody has labelled them. K-Means would give you hard boundaries — each point belongs to exactly one cluster. But what if a data point sits right between two clusters? Shouldn't we express that uncertainty?

Gaussian Mixture Models (GMMs) do exactly this. They model your data as a mixture of Gaussian distributions, assigning each point a probability of belonging to each cluster — not a hard label. By the end of this post, you'll implement a GMM from scratch using the EM algorithm and cluster real data from the Old Faithful geyser.

If you haven't read the EM algorithm tutorial yet, don't worry — we'll explain everything you need here. But if you want the gentle coin-toss introduction first, start there.

Quick Win: Cluster Old Faithful Eruptions

Old Faithful is a geyser in Yellowstone National Park. Its eruptions come in two types: short (~2 min) with short waits, and long (~4.5 min) with long waits. Let's discover these clusters automatically.

import numpy as np
import matplotlib.pyplot as plt

def gaussian_pdf(x, mu, sigma2):
    """Univariate Gaussian probability density function."""
    return (1 / np.sqrt(2 * np.pi * sigma2)) * np.exp(-0.5 * (x - mu)**2 / sigma2)

def fit_gmm(data, K=2, max_iter=50, tol=1e-6):
    """
    Fit a Gaussian Mixture Model using EM.

    Args:
        data: 1D array of observations
        K: number of mixture components
        max_iter: maximum EM iterations
        tol: convergence threshold

    Returns:
        means, variances, weights, responsibilities, log_likelihoods
    """
    N = len(data)

    # Initialise parameters
    rng = np.random.default_rng(42)
    means = rng.choice(data, K, replace=False)
    variances = np.full(K, np.var(data))
    weights = np.full(K, 1.0 / K)

    log_liks = []

    for iteration in range(max_iter):
        # E-STEP: Compute responsibilities
        resp = np.zeros((N, K))
        for k in range(K):
            resp[:, k] = weights[k] * gaussian_pdf(data, means[k], variances[k])
        resp /= resp.sum(axis=1, keepdims=True)

        # M-STEP: Update parameters
        Nk = resp.sum(axis=0)
        means = (resp * data[:, None]).sum(axis=0) / Nk
        variances = (resp * (data[:, None] - means)**2).sum(axis=0) / Nk
        weights = Nk / N

        # Compute log-likelihood
        ll = np.sum(np.log(sum(
            weights[k] * gaussian_pdf(data, means[k], variances[k])
            for k in range(K)
        )))
        log_liks.append(ll)

        if iteration > 0 and abs(log_liks[-1] - log_liks[-2]) < tol:
            break

    return means, variances, weights, resp, log_liks

# Old Faithful eruption durations
eruptions = np.array([
    3.6, 1.8, 3.333, 2.283, 4.533, 2.883, 4.7, 3.6, 1.95, 4.35,
    1.833, 3.917, 4.2, 1.75, 4.7, 2.167, 1.75, 4.8, 1.6, 4.25,
    1.8, 1.75, 3.45, 3.067, 4.533, 3.6, 1.967, 4.083, 3.85, 4.433,
    4.3, 4.467, 3.367, 4.033, 3.833, 2.017, 1.867, 4.833, 1.833, 4.783,
    4.35, 1.883, 4.567, 1.75, 4.533, 3.317, 3.833, 2.1, 4.633, 2.0,
    4.8, 4.716, 1.833, 4.833, 1.733, 4.883, 3.717, 1.667, 4.567, 4.317,
    2.233, 4.5, 1.75, 4.8, 1.817, 4.4, 4.167, 4.7, 2.067, 4.7,
    4.033, 1.967, 4.5, 4.0, 1.983, 5.067, 2.017, 4.567, 3.883, 3.6,
    4.133, 4.333, 4.1, 2.633, 4.067, 4.933, 3.95, 4.517, 2.167, 4.0,
    2.2, 4.333, 1.867, 4.817, 1.833, 4.3, 4.667, 3.75, 1.867, 4.9,
    2.483, 4.367, 2.1, 4.5, 4.05, 1.867, 4.583, 1.883, 4.367, 1.75,
    4.417, 1.967, 4.185, 1.883, 4.3, 1.767, 4.533, 2.35, 4.533, 4.45,
    3.567, 4.5, 4.15, 3.817, 3.917, 4.45, 2.0, 4.283, 4.767, 4.533,
    1.85, 4.25, 1.983, 2.25, 4.75, 4.117, 2.15, 4.417, 1.817, 4.467,
])

means, variances, weights, resp, log_liks = fit_gmm(eruptions)

print(f"Cluster 1: μ = {means[0]:.2f} min, σ = {np.sqrt(variances[0]):.2f}, weight = {weights[0]:.2f}")
print(f"Cluster 2: μ = {means[1]:.2f} min, σ = {np.sqrt(variances[1]):.2f}, weight = {weights[1]:.2f}")

The result: EM discovers two clear groups — short eruptions (~2.0 min, ~35% of data) and long eruptions (~4.3 min, ~65% of data). Let's visualise:

x_plot = np.linspace(1, 6, 300)
mixture = sum(weights[k] * gaussian_pdf(x_plot, means[k], variances[k]) for k in range(2))

plt.figure(figsize=(10, 5))
plt.hist(eruptions, bins=20, density=True, alpha=0.5, color='steelblue', label='Data')
for k in range(2):
    component = weights[k] * gaussian_pdf(x_plot, means[k], variances[k])
    plt.plot(x_plot, component, '--', linewidth=2, label=f'Component {k+1}: μ={means[k]:.2f}')
plt.plot(x_plot, mixture, 'k-', linewidth=2, label='Mixture')
plt.xlabel('Eruption Duration (minutes)')
plt.ylabel('Density')
plt.title('GMM Fit to Old Faithful Eruption Data')
plt.legend()
plt.show()

What Just Happened?

A GMM assumes your data was generated by a process like this:

Randomly pick one of $K$ Gaussian components (with probabilities $\pi_1, \ldots, \pi_K$ )
Draw a data point from that Gaussian

The catch: you observe the data points, but not which component generated each one. This hidden information — the component assignments — is the latent variable that EM handles.

E-Step: "Who Generated This Point?"

Given our current parameter estimates, we compute the responsibility of each component for each data point:

Think of $r_{nk}$ as "how much does component $k$ claim ownership of data point $x_n$ ?" If a point sits right at the mean of component 1 and far from component 2, then $r_{n1} \approx 1$ and $r_{n2} \approx 0$ .

This is exactly the same idea as the soft assignments in the coin toss EM — instead of hard clustering, each point gets fractional membership.

# Let's trace the responsibilities for a few data points
print("Responsibilities (probability of belonging to each cluster):")
print(f"{'Duration':>10s}  {'Cluster 1':>10s}  {'Cluster 2':>10s}  {'Assignment':>12s}")
print("-" * 50)
for idx in [0, 1, 4, 7, 23]:
    dur = eruptions[idx]
    r1, r2 = resp[idx]
    label = "Short" if r1 > 0.5 else "Long"
    print(f"{dur:10.2f}  {r1:10.4f}  {r2:10.4f}  {label:>12s}")

M-Step: "Update the Gaussians"

Using the responsibilities as weights, we re-estimate each component's parameters. The formulas are weighted maximum likelihood estimates — the same MLE you'd compute with complete data, but weighted by the responsibilities:

Effective number of points assigned to component $k$ :

Updated mean (weighted average of data):

Updated variance (weighted variance):

Updated mixing weight:

Notice the pattern: if a point has $r_{nk} = 0.8$ for component $k$ , then 80% of that point's "mass" goes toward updating component $k$ 's parameters. This is exactly how weighted MLE works.

Convergence

Let's verify that EM increases the log-likelihood at every step:

plt.figure(figsize=(8, 4))
plt.plot(range(1, len(log_liks) + 1), log_liks, 'g-o', markersize=6)
plt.xlabel('Iteration')
plt.ylabel('Log-Likelihood')
plt.title('EM Monotonically Increases the Log-Likelihood')
plt.grid(True, alpha=0.3)
plt.show()

Going Deeper

The Full GMM Log-Likelihood

The log-likelihood we're maximising is:

The $\log$ of a sum makes this hard to optimise directly — you can't just take the derivative and set it to zero. That's why we need EM: it transforms this intractable problem into a sequence of tractable weighted MLE problems.

Choosing K: How Many Clusters?

GMMs require you to specify the number of components $K$ . Too few and you underfit; too many and you overfit. Common approaches:

BIC (Bayesian Information Criterion):

where $p$ is the number of parameters ( $3K - 1$ for univariate GMMs: $K$ means + $K$ variances + $K-1$ weights).

def compute_bic(data, K_range=range(1, 6)):
    """Compute BIC for different numbers of components."""
    bics = []
    for K in K_range:
        means, variances, weights, _, log_liks = fit_gmm(data, K=K)
        n_params = 3 * K - 1
        bic = -2 * log_liks[-1] + n_params * np.log(len(data))
        bics.append(bic)
    return list(K_range), bics

K_values, bics = compute_bic(eruptions)
print("BIC scores:")
for k, bic in zip(K_values, bics):
    marker = " ← best" if bic == min(bics) else ""
    print(f"  K={k}: BIC = {bic:.1f}{marker}")

Lower BIC is better. For Old Faithful, $K=2$ should win clearly.

GMM vs K-Means

Feature	K-Means	GMM
Assignments	Hard (0 or 1)	Soft (probabilities)
Cluster shape	Spherical only	Elliptical (with full covariance)
Uncertainty	No	Yes — $r_{nk}$ quantifies it
Output	Cluster labels	Full generative model
Speed	Faster	Slower (but more informative)

K-Means is actually a special case of GMMs where all covariances are equal and small, forcing the responsibilities toward 0 or 1.

Common Pitfalls

Singularities — If a component collapses onto a single data point, its variance goes to zero and the likelihood goes to infinity. Fix: add a small floor to variances (e.g., $\sigma^2 \geq 10^{-6}$ ).
Initialisation sensitivity — Like all EM algorithms, GMMs can get stuck in local optima. Run multiple initialisations and keep the best result.
Label switching — The component labels are arbitrary. After fitting, component 1 in one run might correspond to component 2 in another.
High dimensions — In high dimensions, full covariance matrices have $\mathcal{O}(d^2)$ parameters per component. Use diagonal or tied covariances to reduce complexity.

When NOT to Use GMMs

Non-Gaussian clusters — If your clusters are banana-shaped or have complex geometry, GMMs will struggle. Consider DBSCAN or spectral clustering.
Very high dimensions — The curse of dimensionality makes density estimation unreliable. Reduce dimensions first (PCA, t-SNE).
Huge datasets — Full EM on millions of points is slow. Use mini-batch variants or the scikit-learn implementation.

Multivariate Extension

Our implementation handles 1D data. For $d$ -dimensional data, the Gaussian PDF becomes:

The M-step updates become matrix operations, but the structure is identical: weighted means, weighted covariance matrices, updated mixing weights.

Deep Dive: The Paper

Dempster, Laird, and Rubin (1977)

The EM algorithm for GMMs is a specific instance of the general framework introduced by Dempster, A.P., Laird, N.M. and Rubin, D.B. (1977) "Maximum Likelihood from Incomplete Data via the EM Algorithm", Journal of the Royal Statistical Society, Series B, 39(1), 1–38.

The paper unified a collection of ad-hoc algorithms that had been used in various fields — genetics, astronomy, engineering — under a single theoretical framework. Their key contribution was proving that every EM iteration is guaranteed to increase the likelihood (or leave it unchanged).

"Each iteration of a generalized EM algorithm increases the likelihood and the sequence of likelihood values converges."
— Dempster, Laird & Rubin (1977)

The Complete-Data Framework

In GMM terms, the paper's framework maps to:

Observed (incomplete) data $X$ : the data points $x_1, \ldots, x_N$
Latent variables $Z$ : the component assignments $z_1, \ldots, z_N$ (which Gaussian generated each point)
Complete data: $(X, Z)$ — if we knew both, MLE would be trivial
Parameters $\theta = \{\mu_k, \sigma_k^2, \pi_k\}$ : what we want to estimate

The E-step computes $Q(\theta \mid \theta^{(t)})$ — the expected complete-data log-likelihood:

The M-step maximises this over $\theta$ , yielding the weighted MLE formulas we derived earlier.

Bishop's Modern Treatment

Bishop's Pattern Recognition and Machine Learning (2006), Chapter 9 provides the most accessible modern treatment of GMMs. Key insights include:

GMMs as a density estimator — beyond clustering, GMMs give you a full probability model $p(x)$ that you can sample from, evaluate, and use for anomaly detection
Relationship to K-Means — Bishop shows that K-Means is the zero-temperature limit of GMMs
Bayesian GMMs — placing priors on the parameters avoids singularities and enables automatic model selection. This connects to MCMC methods for posterior inference.

The Stanford EM Lecture

The original code in this tutorial references the Stanford Statistics 315a lecture notes on EM by Trevor Hastie and Rob Tibshirani. Slide 13 gives the GMM update equations we implemented.

Try It Yourself

The interactive notebook includes exercises:

2D clustering — Extend the algorithm to cluster both eruption duration and waiting time simultaneously
Model selection — Plot BIC vs K and find the optimal number of components
Initialisation experiments — Run with 10 different random seeds and compare results
Compare with scikit-learn — Verify your implementation matches sklearn.mixture.GaussianMixture
Anomaly detection — Use the fitted GMM to flag data points with low probability as outliers

Interactive Tools

Distribution Explorer — Visualise Gaussian and other distributions that form the components of mixture models

The EM Algorithm: An Intuitive Guide — The coin-toss introduction to EM that GMMs build upon
Maximum Likelihood Estimation from Scratch — The likelihood foundations that EM maximises in each M-step
MCMC Island Hopping: Understanding Metropolis-Hastings — A Bayesian alternative for when you want posterior distributions over GMM parameters

Frequently Asked Questions

What are responsibilities in a Gaussian Mixture Model?

Responsibilities are the probabilities that each data point belongs to each mixture component. During the E-step, the model computes how likely it is that component k generated data point n, given the current parameter estimates. These soft assignments sum to 1 across all components for each data point.

How does a GMM differ from simply running K-Means?

A GMM provides soft probabilistic assignments rather than hard labels, captures elliptical cluster shapes through full covariance matrices, and outputs a complete generative model you can sample from. K-Means only produces spherical cluster boundaries and gives no measure of uncertainty for points near cluster edges.

Why does the EM algorithm sometimes converge to a poor solution?

EM is guaranteed to increase the log-likelihood at every iteration, but it only finds a local maximum, not necessarily the global one. The final solution depends on the initial parameter values. Running EM multiple times with different random initialisations and selecting the result with the highest log-likelihood is the standard workaround.

What is the singularity problem in GMMs and how do I avoid it?

If a Gaussian component collapses onto a single data point, its variance approaches zero and the likelihood goes to infinity. This creates a degenerate solution. You can prevent this by adding a small floor to the variance (e.g., 1e-6), using regularisation, or placing a Bayesian prior on the covariance parameters.

How do I decide how many mixture components K to use?

The Bayesian Information Criterion (BIC) is the most common approach. Fit GMMs with different values of K, compute the BIC for each, and choose the K with the lowest BIC. The BIC balances model fit (log-likelihood) against complexity (number of parameters), penalising models that are unnecessarily complex.

DEV Community: Berkan Sesen

Gaussian Process Regression: The Bayesian Approach to Curve Fitting

Let's Build It

What Just Happened?

The Kernel: Encoding Smoothness

Building the Covariance Matrices

The Prediction Equations

Going Deeper

GP Prior: What the Kernel Believes Before Seeing Data

Hyperparameter Optimisation: Marginal Log-Likelihood

Validating with sklearn

When NOT to Use Gaussian Processes

Where This Comes From

Rasmussen & Williams (2006)

Ebden (2008) — Our Implementation Reference

Historical Context

Further Reading

Try It Yourself

Interactive Tools

Related Posts

Frequently Asked Questions

What does a Gaussian process actually model?

What does the length-scale hyperparameter control?

Why do Gaussian processes scale poorly to large datasets?

How is GP regression related to Bayesian optimisation?

Can I use kernels other than the RBF?

Hyperparameter Optimization: Grid vs Random vs Bayesian

Quick Win: Run All Three Methods

Setup: The Problem

Method 1: Grid Search

Method 2: Random Search

Method 3: Bayesian Optimization (Gaussian Process)

Results Comparison

What Just Happened?

Grid Search: The Cartesian Product

Random Search: The Bergstra-Bengio Insight

Bayesian Optimization: Learning from History

Going Deeper

The GP Surrogate in Action

GP Predictive Distribution

Acquisition Functions: Deciding Where to Look Next

Exploration vs Exploitation

When NOT to Use Bayesian Optimization

Hyperparameters

Deep Dive: The Papers

Bergstra & Bengio (2012) — Random Search for Hyper-Parameter Optimization

Snoek, Larochelle & Adams (2012) — Practical Bayesian Optimization

Mockus et al. (1978) & Jones et al. (1998) — The Origins

Historical Timeline

Modern Alternatives: Optuna and Hyperopt

Further Reading

Interactive Tools

Related Posts

Try It Yourself

Frequently Asked Questions

What is hyperparameter optimisation and why does it matter?

When should I use random search instead of grid search?

What is Bayesian optimisation and when is it worth the overhead?

How many hyperparameter evaluations should I run?

Should I use cross-validation during hyperparameter optimisation?

Policy Gradients: REINFORCE from Scratch with NumPy

Quick Win: Run the Algorithm

Visualise the Learning

What Just Happened?

The Policy Network (4 → 100 → 1)

The Policy Gradient Signal

The PG Magic Line

Reward Discounting

Reward Standardisation (Variance Reduction)

Manual Backpropagation

Batch Gradient Accumulation

RMSProp from Scratch

Going Deeper

The Policy Gradient Theorem

The Log-Likelihood Trick

Why Reward Standardisation Reduces Variance

Hyperparameters

Value-Based vs Policy-Based: When to Use Each

When NOT to Use Vanilla REINFORCE

Deep Dive: The Papers