DEV Community: Berkan Sesen

Value Iteration vs Q-Learning: Dynamic Programming Meets RL

Berkan Sesen — Mon, 04 May 2026 13:07:14 +0000

You have a map of the frozen lake. Every crack in the ice, every slippery patch, every hole is marked. You can sit at your desk and plan the perfect route before stepping foot on the ice. That is value iteration.

Now imagine you have no map. You lace up your boots and start walking. You slip, you fall into holes, you backtrack. But each time you learn a little more about which moves pay off and which ones do not. That is Q-learning.

Both approaches solve the same problem (finding the best policy in a Markov Decision Process), but they start from radically different assumptions about what you know. In our earlier Q-learning post, we focused purely on the model-free approach. This post puts the two side by side on the same FrozenLake environment, so you can see exactly what a model buys you, and what you give up when you do not have one.

By the end of this post, you will have implemented both value iteration and Q-learning from scratch, compared their convergence and policies head-to-head, and understood the Bellman equation that underpins them both.

Quick Win: Run Both Algorithms

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

Watch value iteration discover optimal state values in just a few sweeps, with "heat" radiating outward from the goal:

Here is the complete implementation for both methods:

import numpy as np
import gymnasium as gym

# ── Value Iteration (model-based) ──────────────────────────
def value_iteration(env, gamma=0.95, theta=1e-8):
    """Compute optimal V* using the Bellman optimality equation."""
    nS = env.observation_space.n
    nA = env.action_space.n
    V = np.zeros(nS)

    while True:
        delta = 0
        for s in range(nS):
            action_values = np.zeros(nA)
            for a in range(nA):
                for prob, next_s, reward, done in env.unwrapped.P[s][a]:
                    action_values[a] += prob * (reward + gamma * V[next_s])
            best_value = np.max(action_values)
            delta = max(delta, abs(best_value - V[s]))
            V[s] = best_value
        if delta < theta:
            break

    # Extract greedy policy from V*
    policy = np.zeros(nS, dtype=int)
    for s in range(nS):
        action_values = np.zeros(nA)
        for a in range(nA):
            for prob, next_s, reward, done in env.unwrapped.P[s][a]:
                action_values[a] += prob * (reward + gamma * V[next_s])
        policy[s] = np.argmax(action_values)

    return V, policy

# ── Q-Learning (model-free) ────────────────────────────────
def q_learning(env, n_episodes=10000, gamma=0.95, lr=0.8,
               epsilon_start=1.0, epsilon_end=0.01, decay_rate=7e-3):
    """Tabular Q-learning with epsilon-greedy exploration."""
    nS = env.observation_space.n
    nA = env.action_space.n
    Q = np.zeros((nS, nA))
    rewards = []
    epsilon = epsilon_start

    for ep in range(n_episodes):
        state, _ = env.reset()
        for t in range(100):
            if np.random.uniform() <= epsilon:
                action = env.action_space.sample()
            else:
                action = np.argmax(Q[state, :])

            next_state, reward, terminated, truncated, _ = env.step(action)
            # Q-learning update
            Q[state, action] += lr * (
                reward + gamma * np.max(Q[next_state, :]) - Q[state, action]
            )
            state = next_state
            if terminated or truncated:
                rewards.append(reward)
                epsilon = epsilon_end + (epsilon_start - epsilon_end) * np.exp(-decay_rate * ep)
                break

    policy = np.argmax(Q, axis=1)
    return Q, policy, rewards

# ── Run both on FrozenLake ─────────────────────────────────
env = gym.make('FrozenLake-v1', is_slippery=True)

V_star, vi_policy = value_iteration(env, gamma=0.95)
Q_star, ql_policy, ql_rewards = q_learning(env, n_episodes=10000, gamma=0.95)

The result: Value iteration converges in 184 sweeps and produces a policy that succeeds ~73% of the time. Q-learning, after 10,000 episodes of trial and error, learns a policy that also achieves ~73% success, and agrees with the VI policy on 14 out of 16 states. Both methods find near-identical strategies, but through very different paths.

What Just Happened?

Both algorithms answer the same question: "What is the best action in every state?" But they go about it in fundamentally different ways.

Value Iteration: Planning with a Blueprint

Value iteration has access to the environment's full transition model $P(s' \mid s, a)$ . This is the complete blueprint: for every state and action, you know exactly which states you might land in and with what probability.

The algorithm sweeps through every state, computing the value of the best action using the Bellman optimality equation:

Each sweep propagates value information one step further from the goal. In the GIF above, you can see this: after sweep 0, only the state next to the goal has any value (0.333). By sweep 5, the values have spread across the grid. By sweep 100, they have stabilised.

The key line in the code is this inner loop:

for prob, next_s, reward, done in env.unwrapped.P[s][a]:
    action_values[a] += prob * (reward + gamma * V[next_s])

This sums over all possible outcomes of taking action $a$ from state $s$ , weighting each by its transition probability. No randomness, no sampling; it is a deterministic computation over the full model.

Q-Learning: Learning by Doing

Q-learning has no access to the transition model. It learns by interacting with the environment, collecting $(s, a, r, s')$ tuples, and updating its Q-table one experience at a time:

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

This is the temporal difference (TD) update. The term $r + \gamma \max_{a'} Q(s', a')$ is the TD target: what the agent thinks the return should be based on the immediate reward plus the estimated future value. The difference between this target and the current estimate $Q(s, a)$ is the TD error, which drives learning.

Because Q-learning relies on sampled experience rather than exhaustive computation, it needs many more interactions (10,000 episodes vs 184 sweeps). It also needs an exploration strategy (epsilon-greedy) to ensure it visits enough state-action pairs to build an accurate Q-table. If you have already read our Q-learning tutorial, these mechanics will be familiar.

Why Both Reach the Same Answer

This is not a coincidence. Both algorithms are solving the same Bellman optimality equation. Value iteration solves it through repeated full sweeps over the state space. Q-learning solves it through stochastic approximation: each sampled experience nudges the Q-values toward the true solution, one step at a time.

Given enough sweeps, value iteration converges exactly. Given enough episodes, Q-learning converges asymptotically (with probability 1, under mild conditions on the learning rate and exploration). On FrozenLake, both methods produce policies that agree on 14 of 16 states and achieve the same ~73% success rate.

Going Deeper

Why Not 100%? The Stochasticity Tax

Even the optimal policy only succeeds about 73% of the time on slippery FrozenLake. This is not a bug in the algorithm. The environment is genuinely stochastic: each action has only a 1/3 chance of going in the intended direction, with 1/3 probability of sliding in each perpendicular direction. Some starting positions are simply doomed to fail because all paths to the goal pass near holes, and the ice will occasionally slide you in.

Convergence: 184 Sweeps vs 10,000 Episodes

Value iteration converges to the exact solution in 184 sweeps:

The Bellman error (maximum change in any state value) decreases exponentially. This is because value iteration is a contraction mapping: each sweep brings V closer to the true V* by a factor of at least $\gamma$ . With $\gamma = 0.95$ , the error shrinks by at least 5% per sweep, guaranteeing convergence.

Q-learning, by contrast, follows a noisier path:

The rolling average hovers around 40-60% during training because the agent is still exploring (epsilon > 0). But the extracted greedy policy, evaluated after training with epsilon = 0, achieves the same 73% as value iteration.

The Model-Based vs Model-Free Tradeoff

This comparison crystallises one of the deepest tradeoffs in reinforcement learning:

Property	Value Iteration	Q-Learning
Needs transition model?	Yes (env.P)	No
Steps to converge	184 sweeps	~10,000 episodes
Optimality guarantee	Exact	Asymptotic
Works for unknown environments?	No	Yes
Memory	O(\|S\|)	O(\|S\| × \|A\|)

Value iteration is faster and guarantees exact optimality, but it requires something that is rarely available in practice: the full transition model $P(s' \mid s, a)$ . In robotics, game-playing, or any complex real-world task, you almost never have this. That is why model-free methods like Q-learning (and its deep successor, DQN) dominate modern RL.

Hyperparameter Sensitivity

The original code uses a high learning rate ( $\alpha = 0.8$ ) and fast epsilon decay ( $\text{decay\_rate} = 7 \times 10^{-3}$ ). This means Q-learning explores aggressively early on and then commits to exploitation within about 1,000 episodes. The high learning rate works here because FrozenLake has a small, discrete state space. For larger problems, you would need to lower $\alpha$ considerably (our DQN post uses 0.001 with a neural network).

Value iteration, by contrast, has no learning rate. The discount factor $\gamma = 0.95$ is the only tunable parameter, and it has a clear interpretation: how much to value future rewards relative to immediate ones. Higher gamma means the agent plans further ahead but converges more slowly.

When to Use Which

Use value iteration when:

You have a complete model of the environment (transition probabilities and rewards)
The state space is small enough to sweep over exhaustively
You need guaranteed optimality

Use Q-learning when:

You can only interact with the environment through trial and error
The model is unknown or too complex to specify
You are willing to trade computation for generality

In practice, most interesting problems fall into the Q-learning camp, which is why model-free methods get so much attention. Not all model-free approaches use value functions, though. Methods like the cross-entropy method and simulated annealing search policy space directly without ever estimating state values. But understanding value iteration is essential because it reveals the Bellman equation that underlies all value-based RL. As we saw in our policy gradient post, even gradient-based methods ultimately try to maximise the same value function.

Deep Dive: The Papers

Bellman's Foundation

Value iteration traces directly to Richard Bellman's 1957 monograph Dynamic Programming. Bellman introduced the principle of optimality:

"An optimal policy has the property that whatever the initial state and initial decision are, the remaining decisions must constitute an optimal policy with regard to the state resulting from the first decision."

This recursive insight leads to the Bellman optimality equation. For the state-value function:

And for the action-value function (the one Q-learning estimates):

Value iteration simply applies the first equation as an update rule, sweeping over all states until convergence. The convergence is guaranteed because the Bellman operator is a contraction in the sup-norm with coefficient $\gamma < 1$ (proven by Bellman himself and later formalised by Denardo, 1967).

Watkins' Q-Learning

Q-learning was introduced by Christopher Watkins in his 1989 PhD thesis at Cambridge, with the convergence proof published in Watkins & Dayan (1992). The key insight was that you can learn $Q^*$ directly from experience, without ever knowing the transition model:

Watkins & Dayan proved that Q-learning converges to $Q^*$ with probability 1, provided:

All state-action pairs are visited infinitely often
The learning rate $\alpha$ satisfies: $\sum \alpha_t = \infty$ and $\sum \alpha_t^2 < \infty$

The first condition is why we need epsilon-greedy exploration. The second is a standard stochastic approximation requirement (Robbins-Monro conditions). In practice, we use a fixed or slowly decaying learning rate and rely on the algorithm converging "well enough" rather than proving formal convergence.

The DP-RL Connection

Sutton & Barto's Reinforcement Learning: An Introduction (2nd ed., 2018) makes the connection explicit in Chapters 4 and 6. Value iteration is presented as a dynamic programming method (Chapter 4), while Q-learning is a temporal difference method (Chapter 6). The book shows that TD methods can be viewed as sampling-based approximations to DP: where DP backs up values using the full distribution over successors, TD methods back up using a single sampled successor.

This connection runs deep. Every model-free RL algorithm, from Q-learning to DQN to policy gradients, is implicitly solving a Bellman equation. The difference is in how they approximate the expectation: through tabular sweeps (DP), sampled transitions (TD), or complete episode returns (Monte Carlo).

Try It Yourself

The interactive notebook includes exercises:

Non-slippery mode: Set is_slippery=False and compare. Both methods should now achieve ~100% success. How does this change the convergence speed?
8x8 grid: Try FrozenLake8x8-v1. Value iteration still works perfectly. How does Q-learning cope with the larger state space?
Learned transition model: The original code includes a learn_trans_matrix() function that estimates $P(s' \mid s, a)$ from random play, then runs VI on the learned model. Try this hybrid approach. How many random episodes do you need before the learned model matches the true one?
Discount factor sensitivity: Vary $\gamma$ from 0.5 to 0.99 and plot the success rate for both methods. When does a low gamma hurt?

Understanding value iteration gives you the theoretical bedrock of RL. Understanding Q-learning gives you the practical tool that works when models are not available. Together, they frame the central tradeoff that drives all of modern reinforcement learning.

Interactive Tools

Q-Learning Visualiser — Watch value iteration and Q-learning converge on grid worlds in the browser

Q-Learning on Frozen Lake from Scratch — Deep dive into tabular Q-learning on the same environment
Deep Q-Networks: Experience Replay and Target Networks — Scaling Q-learning with neural networks
Policy Gradients: REINFORCE from Scratch — The policy-based alternative to value methods
Cross-Entropy Method: Evolution-Style RL — A gradient-free approach to the same control problems

Frequently Asked Questions

What is the difference between value iteration and Q-learning?

Value iteration is a dynamic programming method that requires a complete model of the environment (transition probabilities and rewards) and sweeps through all states systematically. Q-learning is model-free: it learns from experience without knowing the environment dynamics. Both converge to optimal values, but value iteration is faster when a model is available, while Q-learning works when it is not.

What is the Bellman equation?

The Bellman equation expresses the value of a state as the immediate reward plus the discounted value of the next state. It is the foundation of both value iteration and Q-learning. Value iteration solves it by iterating the equation across all states until convergence. Q-learning solves it incrementally by updating one state-action pair at a time from experience.

When should I use dynamic programming instead of Q-learning?

Use dynamic programming (value iteration, policy iteration) when you have a complete and accurate model of the environment. This is common in games with known rules, inventory management, and operations research. When the model is unknown, too complex, or too large to enumerate, use model-free methods like Q-learning.

What is the difference between value iteration and policy iteration?

Value iteration updates the value function using the Bellman optimality equation until convergence, then extracts the policy. Policy iteration alternates between evaluating the current policy exactly and improving it greedily. Policy iteration often converges in fewer iterations but each iteration is more expensive. For small state spaces, both work well.

Does value iteration always converge?

Yes, for finite MDPs with a discount factor less than 1. The Bellman operator is a contraction mapping, guaranteeing convergence at a geometric rate. The number of iterations needed depends on the discount factor (higher gamma means slower convergence) and the desired precision. In practice, convergence is usually fast for small to moderate state spaces.

Custom Likelihoods in PyMC: One-Inflated Beta Regression for Loan Repayment

Berkan Sesen — Fri, 01 May 2026 08:47:52 +0000

When a borrower takes out a personal loan, they might repay every penny, default entirely, or land anywhere in between. The interesting variable is the fraction eventually recovered: a number between 0 and 1 for each loan in the portfolio. Plot the distribution across thousands of loans and it looks like a smooth Beta curve with a tall spike bolted on at the right edge — a mass of borrowers who repaid in full.

That spike is good news for the lender, but a headache for the modeller. Standard Beta regression handles continuous outcomes on (0, 1), but it cannot produce a point mass at the boundary. Logistic regression predicts a binary paid-or-not label, throwing away the partial repayment information. Neither tool fits the data you actually have.

In our first PyMC post, we built hierarchical models using built-in distributions. In the second, we handled non-standard likelihoods with pm.Potential for right-censored survival data.

This post takes the final step: writing a piecewise log-likelihood from scratch for a mixture of continuous and discrete components. By the end, you will construct a One-Inflated Beta (OIB) regression in PyMC, hand-code the Beta log-density, and infer how borrower characteristics drive both the probability of full repayment and the expected partial repayment fraction.

Let's Build It

Click the badge below to open the full interactive notebook:

We will generate synthetic loan data for 2,000 borrowers, fit an OIB regression model, and recover the true data-generating parameters.

import numpy as np
import pymc as pm
import pytensor.tensor as pt
import arviz as az
import matplotlib.pyplot as plt

np.random.seed(42)

# --- Generate synthetic loan data ---
N = 2000
credit_score = np.random.normal(0, 1, N)
loan_to_value = np.random.normal(0, 1, N)
interest_rate = np.random.normal(0, 1, N)
income_ratio = np.random.normal(0, 1, N)
X = np.column_stack([credit_score, loan_to_value, interest_rate, income_ratio])
feature_names = ['credit_score', 'loan_to_value', 'interest_rate', 'income_ratio']

# True parameters
true_psi = np.array([0.5, 0.8, -0.6, -0.4, 0.3])    # pi coefficients
true_delta = np.array([0.3, 0.5, -0.3, -0.2, 0.2])   # theta coefficients
true_phi = 5.0                                          # Beta precision

# Per-loan probability of full repayment (logistic link)
logit_pi = true_psi[0] + X @ true_psi[1:]
pi_true = 1 / (1 + np.exp(-logit_pi))

# Per-loan mean partial repayment (logistic link)
logit_theta = true_delta[0] + X @ true_delta[1:]
theta_true = 1 / (1 + np.exp(-logit_theta))

# Beta shape parameters from mean-precision
alpha_true = theta_true * true_phi
beta_true = (1 - theta_true) * true_phi

# Sample from the OIB mixture
u = np.random.uniform(size=N)
repayment = np.where(u < pi_true, 1.0, np.random.beta(alpha_true, beta_true))

n_full = (repayment == 1.0).sum()
print(f"Fully repaid: {n_full}/{N} ({n_full/N:.1%})")
print(f"Partial repayment: {N - n_full}/{N} ({(N - n_full)/N:.1%})")

Of 2,000 loans, 1,214 (60.7%) are fully repaid and 786 (39.3%) show partial repayment. The histogram immediately reveals the two populations: a tall spike at 1.0 and a continuous spread below it. No single standard distribution can capture both. Now let's build the OIB model.

# Split observations by type
full_idx = np.where(repayment == 1.0)[0]
partial_idx = np.where(repayment < 1.0)[0]
partial_values = repayment[partial_idx]

with pm.Model() as oib_model:
    # Pi sub-model: probability of full repayment (logistic link)
    psi_intercept = pm.Normal('psi_intercept', mu=0, sigma=5)
    psi_coeffs = pm.Normal('psi_coeffs', mu=0, sigma=1, shape=4)
    logit_pi = psi_intercept + pt.dot(X, psi_coeffs)
    pi = pm.Deterministic('pi', pm.math.invlogit(logit_pi))

    # Theta sub-model: mean of partial repayment Beta (logistic link)
    delta_intercept = pm.Normal('delta_intercept', mu=0, sigma=5)
    delta_coeffs = pm.Normal('delta_coeffs', mu=0, sigma=1, shape=4)
    logit_theta = delta_intercept + pt.dot(X, delta_coeffs)
    theta = pm.Deterministic('theta', pm.math.invlogit(logit_theta))

    # Phi: Beta precision (shared across all loans)
    phi = pm.Gamma('phi', alpha=2, beta=0.5)

    # Convert mean-precision to standard Beta parameters
    a = theta * phi
    b = (1 - theta) * phi

    # Expected repayment: E[Y] = pi + (1 - pi) * theta
    E_f = pm.Deterministic('E_f', pi + (1 - pi) * theta)

    # --- Piecewise log-likelihood via pm.Potential ---
    # Fully repaid loans: log(pi_i)
    pm.Potential('ll_full', pt.sum(pt.log(pi[full_idx])))

    # Partial repayments: log(1 - pi_i) + log Beta(y_i | a_i, b_i)
    pa, pb = a[partial_idx], b[partial_idx]
    beta_logp = (pt.gammaln(pa + pb) - pt.gammaln(pa) - pt.gammaln(pb)
                 + (pa - 1) * pt.log(partial_values)
                 + (pb - 1) * pt.log(1 - partial_values))
    pm.Potential('ll_partial', pt.sum(pt.log(1 - pi[partial_idx]) + beta_logp))

with oib_model:
    trace = pm.sample(draws=1000, tune=3000, chains=4,
                      target_accept=0.95, random_seed=42,
                      init='jitter+adapt_diag')

print(az.summary(trace, var_names=['psi_intercept', 'psi_coeffs',
                                    'delta_intercept', 'delta_coeffs', 'phi']))

The trace plots show healthy chains: zero divergences, good mixing across all four chains, and unimodal posteriors centred near the true parameter values. Sampling 4,000 draws per chain with the Potential-based likelihood took about 6 seconds.

You just fitted a custom Bayesian mixture model with 11 free parameters. Now let's understand how each piece works.

What Just Happened?

Two populations, one model

Our data contains two distinct groups. Some borrowers repay their loan in full (repayment fraction = 1.0), and others repay partially (0 < fraction < 1). The OIB model treats this as a mixture: with probability $\pi_i$ the outcome is exactly 1, and with probability $1 - \pi_i$ it follows a Beta distribution.

Both $\pi_i$ and the Beta mean $\theta_i$ vary across borrowers. A high credit score might increase both the chance of full repayment and the expected partial repayment. The model captures these relationships through separate linear predictors with logistic links, ensuring both quantities stay between 0 and 1.

The piecewise log-likelihood

The OIB density is a mixture of a point mass and a continuous distribution. For observation $y_i$ :

Taking logs:

The addition in the second branch is critical: it corresponds to multiplying the mixing weight $(1 - \pi_i)$ by the Beta density in probability space. A common mistake is to write multiplication of two log quantities (i.e. log(1-pi) * log(Beta(...))) instead of addition. That would have no probabilistic interpretation.

We implement this by splitting observations into two groups and adding each group's log-likelihood as a separate pm.Potential:

# Fully repaid: sum of log(pi_i) over fully-repaid loans
pm.Potential('ll_full', pt.sum(pt.log(pi[full_idx])))

# Partial: sum of log(1 - pi_i) + Beta_logpdf(y_i) over partial loans
pm.Potential('ll_partial', pt.sum(pt.log(1 - pi[partial_idx]) + beta_logp))

This pattern should feel familiar from Post 21, where we used pm.Potential to handle right-censored observations. The principle is the same: when your likelihood has distinct branches for different observation types, split them into separate Potential terms.

Hand-coding the Beta log-density

Rather than relying on pm.logp(pm.Beta.dist(...), value), we compute the Beta log-density directly using the gamma function:

beta_logp = (pt.gammaln(a + b) - pt.gammaln(a) - pt.gammaln(b)
             + (a - 1) * pt.log(y) + (b - 1) * pt.log(1 - y))

This follows from the Beta density formula $f(y \mid \alpha, \beta) = \frac{\Gamma(\alpha + \beta)}{\Gamma(\alpha)\Gamma(\beta)} y^{\alpha-1}(1-y)^{\beta-1}$ . Writing it out explicitly has two advantages: you can see exactly what the sampler is differentiating through, and you avoid potential issues with PyMC's internal distribution objects when used inside Potential expressions.

The model structure

The model has three sub-components connected by link functions:

Pi sub-model controls which mixture component generates each observation. A logistic link maps the linear predictor $\psi_0 + \psi_1 x_{\text{credit}} + \psi_2 x_{\text{ltv}} + \psi_3 x_{\text{rate}} + \psi_4 x_{\text{income}}$ to a probability. Positive $\psi_1$ means higher credit scores increase the chance of full repayment.

Theta sub-model sets the mean of the Beta distribution for partial repayments, also through a logistic link with its own coefficients $\delta_0, \ldots, \delta_4$ . This captures a subtlety that pure classification misses: among borrowers who do not fully repay, some covariates still push the partial fraction higher.

Phi is a single shared precision parameter for the Beta component. Higher phi means less variance in partial repayments. It uses a $\text{Gamma}(2, 0.5)$ prior with mean 4, which favours moderate precision values.

Checking the fit

Let's compare the estimated coefficients to the true values we used to generate the data.

summary = az.summary(trace, var_names=['psi_intercept', 'psi_coeffs',
                                        'delta_intercept', 'delta_coeffs', 'phi'])

true_vals = np.concatenate([true_psi, true_delta, [true_phi]])
param_names = (['psi_intercept'] + [f'psi_coeffs[{i}]' for i in range(4)] +
               ['delta_intercept'] + [f'delta_coeffs[{i}]' for i in range(4)] + ['phi'])

fig, ax = plt.subplots(figsize=(8, 6))
y_pos = np.arange(len(param_names))
means = summary['mean'].values
hdi_low = summary['hdi_3%'].values
hdi_high = summary['hdi_97%'].values

ax.errorbar(means, y_pos, xerr=[means - hdi_low, hdi_high - means],
            fmt='o', color='steelblue', capsize=4, label='Posterior (94% HDI)')
ax.scatter(true_vals, y_pos, marker='x', color='crimson', s=80,
           zorder=5, label='True value')
ax.set_yticks(y_pos)
ax.set_yticklabels(param_names)
ax.axvline(0, color='gray', linestyle='--', alpha=0.5)
ax.set_xlabel('Value')
ax.legend(loc='lower right')
ax.set_title('Parameter Recovery: Posterior vs True Values')
plt.tight_layout()

Every true value falls within its 94% highest density interval. The model correctly identifies that credit score has the strongest positive effect on full repayment (psi_coeffs[0] = 0.85, true: 0.8), while loan-to-value ratio is the strongest negative predictor (psi_coeffs[1] = -0.58, true: -0.6). The precision parameter phi is recovered at 5.47 (true: 5.0), and the effective sample sizes all exceed 2,500.

Posterior predictive check

The ultimate test: can the model reproduce the observed data distribution, including the spike at 1.0? Since we used pm.Potential rather than an observed distribution, we generate predictive samples manually from the posterior:

# Extract posterior samples
psi_int_post = trace.posterior['psi_intercept'].values.flatten()
psi_coeff_post = trace.posterior['psi_coeffs'].values.reshape(-1, 4)
delta_int_post = trace.posterior['delta_intercept'].values.flatten()
delta_coeff_post = trace.posterior['delta_coeffs'].values.reshape(-1, 4)
phi_post = trace.posterior['phi'].values.flatten()

rng = np.random.default_rng(42)
n_draws = 500
ppc_samples = np.zeros((n_draws, N))

for i in range(n_draws):
    lp = psi_int_post[i] + X @ psi_coeff_post[i]
    pi_i = 1 / (1 + np.exp(-lp))
    lt = delta_int_post[i] + X @ delta_coeff_post[i]
    theta_i = 1 / (1 + np.exp(-lt))
    a_i, b_i = theta_i * phi_post[i], (1 - theta_i) * phi_post[i]
    u_i = rng.uniform(size=N)
    ppc_samples[i] = np.where(u_i < pi_i, 1.0, rng.beta(a_i, b_i))

fig, ax = plt.subplots(figsize=(8, 4))
ax.hist(repayment, bins=50, density=True, alpha=0.5, color='steelblue',
        label='Observed', edgecolor='white')
ax.hist(ppc_samples.flatten(), bins=50, density=True, alpha=0.4, color='coral',
        label='Posterior predictive', edgecolor='white')
ax.set_xlabel('Repayment Fraction')
ax.set_ylabel('Density')
ax.legend()
ax.set_title('Posterior Predictive Check')
plt.tight_layout()

The posterior predictive distribution matches both the spike at 1.0 and the shape of the partial repayment component. This is something neither pure Beta regression nor logistic regression can achieve.

Going Deeper

The mean-precision parameterisation

The standard Beta distribution uses shape parameters $\alpha$ and $\beta$ , but these are difficult to interpret. A borrower with $\alpha = 2.8$ and $\beta = 2.1$ tells you almost nothing at a glance. The mean-precision reparameterisation solves this:

Now $\mu$ is the mean of the distribution (the expected partial repayment fraction) and $\phi$ is the precision (higher means less spread). The inverse mapping recovers the standard parameters:

In our model, $\mu$ is called $\theta$ and depends on covariates through a logistic link. The precision $\phi$ is shared across all observations, which assumes that the variance of partial repayments (given the mean) is the same for all borrowers. This is a simplification; a fully heteroscedastic model would give $\phi$ its own linear predictor.

Why logistic links?

Both $\pi$ (probability of full repayment) and $\theta$ (mean of the Beta) must live in (0, 1). The logistic function $\sigma(x) = 1 / (1 + e^{-x})$ maps any real-valued linear predictor to this interval. This is the same link function used in logistic regression and Bayesian classification.

The priors reflect the link: $\text{Normal}(0, 5)$ on the intercepts allows the baseline probability to range widely, while $\text{Normal}(0, 1)$ on the slope coefficients gently regularises each covariate's effect. On the logistic scale, a coefficient of 1.0 roughly doubles the odds, so a $\text{Normal}(0, 1)$ prior is mildly informative.

The expected value formula

The overall expected repayment for borrower $i$ combines both components:

This is the E_f deterministic in our model. It allows you to rank borrowers by expected repayment even when their risk profiles differ in how they fail: one borrower might have a high chance of full repayment but low partial repayment if they default, while another has a moderate chance of full repayment but high partial recovery.

Why pm.Potential and not pm.CustomDist?

PyMC offers two ways to implement custom likelihoods. pm.CustomDist lets you define a distribution from its logp function, which would look like this for OIB:

def oib_logp(value, pi, a, b):
    return pt.switch(
        pt.eq(value, 1.0),
        pt.log(pi),
        pt.log(1 - pi) + pm.logp(pm.Beta.dist(alpha=a, beta=b), value)
    )

This is elegant but fragile. The pt.switch operator evaluates both branches for every observation during automatic differentiation.

When value = 1.0, the Beta branch computes pm.logp(Beta, 1.0), which returns negative infinity (since the Beta density is zero at boundaries for $\beta > 1$ ). Even though the switch selects the other branch, the gradient through the infinite branch corrupts the NUTS sampler. The result: 100% divergence rate.

The pm.Potential approach avoids this entirely. By pre-splitting observations into fully-repaid and partial groups, the Beta density is never evaluated at the boundary. This is the same pattern we used for censored data in survival analysis: separate the observation types, compute each group's log-likelihood independently, and add them as Potential terms.

The trade-off is that pm.Potential does not enable pm.sample_posterior_predictive out of the box (you need to write manual prediction code, as we did). For many production workflows, that is a minor inconvenience compared to the reliability gain.

Sampling considerations

Our sampling configuration follows the original code that inspired this tutorial:

3,000 tuning steps with 1,000 posterior draws per chain. The long warm-up helps the NUTS sampler adapt its step size to the geometry of the piecewise likelihood.
4 chains for convergence diagnostics. With $\hat{R}$ and effective sample size, four chains provide reliable evidence that the sampler has explored the full posterior.
target_accept=0.95 raises the acceptance threshold from the default 0.8, which reduces divergences in models with sharp likelihood boundaries.
init='jitter+adapt_diag' initialises each chain near the prior mean with small random perturbations. A practical note from the original code: if covariates have very different scales (e.g., one ranges from 0 to 1 while another ranges from 0 to 200), the default jitter of roughly $\pm 1$ can push initial coefficient values far from reasonable territory. Standardising covariates beforehand, as we did, avoids this.

When to use something else

The OIB model assumes that exactly-one observations arise from a fundamentally different process than partial observations. If instead you have data with a spike at zero (e.g., insurance claims where most customers file nothing), you want a zero-inflated model. If you have spikes at both boundaries, you need a zero-and-one-inflated Beta (ZOIB).

For data with no boundary spikes at all, standard Beta regression (via MLE or Bayesian inference) is simpler and sufficient. The extra complexity of the OIB mixture is only justified when the data genuinely contains a discrete mass at the boundary.

Where This Comes From

The OIB model sits at the intersection of two lines of research: Beta regression for bounded continuous data, and inflated distributions for boundary spikes.

Beta regression: Ferrari and Cribari-Neto (2004)

The foundation is the Beta regression model introduced by Silvia Ferrari and Francisco Cribari-Neto in their 2004 paper "Beta Regression for Modelling Rates and Proportions" (Journal of Applied Statistics, 27(7), 799-815). They observed that rates, proportions, and fractions appear everywhere in applied statistics, yet researchers typically transform them (logit, arcsine) and apply linear regression. This is problematic because the transformation distorts the error structure and complicates interpretation.

Their key insight was to model the response directly as Beta-distributed, using the mean-precision parameterisation we adopted:

where $0 < y < 1$ , $0 < \mu < 1$ is the mean, and $\phi > 0$ is the precision. Ferrari and Cribari-Neto showed that this is a natural exponential family model when parameterised through $\mu$ , and proposed a logit link for the mean:

"The proposed model is useful for situations where the variable of interest is continuous and restricted to the interval (0, 1). [...] A convenient parameterisation of the beta density in terms of the mean and a precision parameter is used."

Their framework supports maximum likelihood estimation, but the Bayesian extension (which we use) adds uncertainty quantification and regularisation through priors. The connection to MLE is direct: the posterior mode of our model with flat priors equals the MLE of Ferrari and Cribari-Neto's model.

Inflated models: Ospina and Ferrari (2010)

The standard Beta has support on the open interval (0, 1), so it cannot assign positive probability to the boundaries 0 or 1. Raydonal Ospina and Silvia Ferrari addressed this in "Inflated Beta Distributions" (Statistical Papers, 51(1), 111-126, 2010). They defined a class of mixed continuous-discrete distributions:

This is exactly the piecewise density we implemented with pm.Potential. The parameter $\pi$ controls the inflation: the probability of observing the boundary value. Ospina and Ferrari also developed zero-inflated and zero-and-one-inflated variants for different boundary patterns.

"In many practical situations, the variable of interest is continuous in the open standard unit interval but may also assume the extreme values zero and/or one with positive probabilities. [...] We introduce a class of inflated beta distributions."

Their work established the theoretical properties (moments, maximum likelihood estimation, score functions) that underpin our Bayesian implementation.

From MLE to MCMC

The original MLE approach estimates $\pi$ , $\mu$ , and $\phi$ by maximising the log-likelihood. The Bayesian version replaces optimisation with MCMC sampling, yielding full posterior distributions rather than point estimates. This is particularly valuable for the OIB model because the piecewise likelihood creates a posterior geometry that point estimates cannot capture: the uncertainty in $\pi$ and $\theta$ is correlated, and the posterior for $\phi$ is often skewed.

Where Ferrari and Cribari-Neto derived score functions by hand, we supply the log-density components to PyMC and let the NUTS sampler handle the rest. The automatic differentiation in PyTensor computes gradients through the gammaln and log operations, enabling efficient Hamiltonian Monte Carlo.

Algorithm summary

The complete OIB regression procedure:

For each observation $i$ , compute $\pi_i = \sigma(\psi_0 + \mathbf{x}_i^\top \boldsymbol{\psi})$ (full repayment probability)
Compute $\theta_i = \sigma(\delta_0 + \mathbf{x}_i^\top \boldsymbol{\delta})$ (partial repayment mean)
Compute Beta shape parameters: $\alpha_i = \theta_i \phi$ , $\beta_i = (1 - \theta_i) \phi$
Evaluate the piecewise log-likelihood: $\log \pi_i$ if $y_i = 1$ , else $\log(1 - \pi_i) + \log \text{Beta}(y_i \mid \alpha_i, \beta_i)$
Sum across all observations and sample the posterior via NUTS

Interactive Tools

Distribution Explorer — Visualise the Beta distribution and other families used in this model
Bayes' Theorem Calculator — Explore Bayesian reasoning interactively

Hierarchical Bayesian Regression with PyMC: When Groups Share Strength — Partial pooling and group-level priors in PyMC
Bayesian Survival Analysis with PyMC: Modelling Customer Churn — Another custom likelihood built in PyMC
From MLE to Bayesian Inference — Why we use priors and posteriors instead of point estimates
MCMC Metropolis-Hastings: An Island-Hopping Guide — The sampling engine behind PyMC

Frequently Asked Questions

When should I use a One-Inflated Beta model instead of logistic regression?

Use OIB when your outcome is a fraction between 0 and 1 with a spike at the boundary value of 1. Logistic regression discards the partial repayment information by collapsing everything into a binary label. OIB preserves both the probability of full repayment and the distribution of partial repayments, giving you richer predictions.

Why use pm.Potential instead of pm.CustomDist for the likelihood?

The pm.CustomDist approach evaluates both branches of the piecewise likelihood for every observation during automatic differentiation. When the Beta density is evaluated at the boundary value of 1.0, it returns negative infinity, which corrupts the NUTS sampler gradients and causes 100% divergences. Splitting observations with pm.Potential avoids evaluating the Beta density at the boundary entirely.

What is the mean-precision parameterisation of the Beta distribution?

Instead of the standard shape parameters alpha and beta, the mean-precision form uses mu (the mean, between 0 and 1) and phi (the precision, controlling spread). This is more interpretable: mu directly tells you the expected partial repayment fraction, while phi tells you how concentrated the distribution is around that mean. The standard parameters are recovered as alpha = mu * phi and beta = (1 - mu) * phi.

How do I check whether the OIB model fits my data well?

Generate posterior predictive samples by drawing from the fitted model and comparing the resulting distribution to the observed data. The key check is whether the model reproduces both the spike at 1.0 (the proportion of fully repaid loans) and the shape of the continuous partial repayment distribution. If either component is mismatched, the model needs adjustment.

Can this model handle spikes at both 0 and 1?

Yes, but you would need a Zero-and-One-Inflated Beta (ZOIB) model. This adds a third mixture component for the spike at zero, with its own probability parameter. The piecewise likelihood gains a third branch, but the pm.Potential implementation pattern remains the same: split observations into three groups and add each group's log-likelihood separately.

Bayesian Survival Analysis with PyMC: Modelling Customer Churn

Berkan Sesen — Wed, 29 Apr 2026 09:53:05 +0000

Every subscription business lives or dies by churn. Whether it is a B2B SaaS platform tracking annual contracts or a consumer app watching monthly renewals, the question is the same: how long will this customer stay? The data seems straightforward. Some subscribers cancelled after a month, others after a year. But a large share of customers are still active. They have not churned yet, and you do not know when, or whether, they will.

A colleague suggested dropping them from the analysis. That felt wrong, and it is: ignoring active customers biases your model toward shorter lifetimes, because you only learn from the people who already left.

The problem has a name: right-censoring. An active customer who signed up 8 months ago tells you something valuable: they survived at least 8 months. You don't know when (or whether) they'll churn, but that lower bound is real information.

Survival analysis handles censoring properly. In our previous post, we built hierarchical models in PyMC for grouped regression. This post extends that toolkit with a new ingredient: the ability to learn from incomplete observations.

By the end, you'll build a Bayesian accelerated failure time (AFT) model in PyMC, handle right-censored data with pm.Potential, compare Weibull and Log-Logistic distributions, and plot individual survival curves for different customer profiles.

Let's Build It

First, let's see the model in action. Click the badge below to open the full interactive notebook:

We'll generate synthetic churn data for 1,000 customers, fit a Weibull AFT model, and plot survival curves.

import numpy as np
import pandas as pd
import pymc as pm
import pytensor.tensor as pt
import arviz as az
import matplotlib.pyplot as plt

np.random.seed(42)

# Generate synthetic churn data: 1,000 customers observed over 24 months
N = 1000
monthly_spend = np.random.normal(100, 30, N).clip(20, 250)
support_tickets = np.random.poisson(3, N).astype(float)

# Standardise covariates
spend_std = (monthly_spend - 100) / 30
tickets_std = (support_tickets - 3) / 2

# True AFT parameters (Gumbel / log-Weibull parameterisation)
true_alpha = np.array([2.5, 0.4, -0.3])  # intercept, spend, tickets
true_s = 0.6

# True log-time: Y = eta + s * W, where W ~ Gumbel(0,1)
eta_true = true_alpha[0] + true_alpha[1] * spend_std + true_alpha[2] * tickets_std
log_time_true = eta_true + true_s * np.random.gumbel(0, 1, N)
time_true = np.exp(log_time_true)

# Administrative censoring at 24 months
observation_window = 24.0
observed_time = np.minimum(time_true, observation_window)
censored = time_true > observation_window  # True = still active
log_observed_time = np.log(observed_time)

print(f"Total customers: {N}")
print(f"Churned: {(~censored).sum()} ({(~censored).mean():.0%})")
print(f"Still active (censored): {censored.sum()} ({censored.mean():.0%})")

Total customers: 1000
Churned: 664 (66%)
Still active (censored): 336 (34%)

Before fitting the Bayesian model, let's look at the empirical survival curve using the Kaplan-Meier estimator. This non-parametric method handles censoring correctly by adjusting the risk set at each event time:

# Kaplan-Meier estimator (manual, no extra dependencies)
order = np.argsort(observed_time)
times_sorted = observed_time[order]
events_sorted = (~censored)[order].astype(int)

km_times = [0.0]
km_survival = [1.0]
n_at_risk = N

for t, event in zip(times_sorted, events_sorted):
    if event:
        km_survival.append(km_survival[-1] * (1 - 1 / n_at_risk))
        km_times.append(t)
    n_at_risk -= 1

fig, ax = plt.subplots(figsize=(8, 4))
ax.step(km_times, km_survival, where='post', color='#2196F3', lw=2)
ax.set_xlabel('Months since signup')
ax.set_ylabel('Survival probability')
ax.set_title('Kaplan-Meier Survival Curve')
ax.set_xlim(0, 25)
ax.set_ylim(0, 1.05)
plt.tight_layout()

Now let's fit the Weibull AFT model. The key insight: if $T \sim \text{Weibull}$ , then $Y = \log T$ follows a Gumbel distribution. So we model log-time with a Gumbel likelihood, which lets us write the linear predictor naturally:

def gumbel_log_sf(y, mu, sigma):
    """Log survival function of the Gumbel distribution."""
    return pt.log1p(-pt.exp(-pt.exp(-(y - mu) / sigma)))

with pm.Model() as weibull_aft:
    # Location coefficients (priors match original code: Normal(0, 2))
    alpha = pm.Normal('alpha', mu=0, sigma=2, shape=3)

    # Scale parameter (must be positive)
    log_s = pm.Normal('log_s', mu=0, sigma=1)
    s = pm.Deterministic('s', pm.math.exp(log_s))

    # Linear predictor for log-time
    eta = alpha[0] + alpha[1] * spend_std + alpha[2] * tickets_std

    # Uncensored customers: standard Gumbel likelihood
    y_obs = pm.Gumbel('y_obs', mu=eta[~censored], beta=s,
                       observed=log_observed_time[~censored])

    # Censored customers: survival function via pm.Potential
    y_cens = pm.Potential('y_cens',
        gumbel_log_sf(log_observed_time[censored], eta[censored], s))

    # Sample the posterior
    trace = pm.sample(1000, tune=2000, cores=4, chains=4,
                      random_seed=42, target_accept=0.9)

print(az.summary(trace, var_names=['alpha', 's']))

You just fit a Bayesian survival model that properly handles censored customers. The alpha coefficients tell you how each covariate affects time-to-churn: positive means longer survival, negative means faster churn. And unlike a point estimate from maximum likelihood, you get full posterior distributions over every parameter.

What Just Happened?

Right-Censoring: Learning from Incomplete Data

The 336 active customers in our data didn't churn during the 24-month observation window. For each one, we know they survived at least 24 months, but not how much longer they'll stay. This is right-censoring: the true event time is somewhere to the right of what we observed.

Standard regression would force you to either drop censored customers (biasing estimates downward) or code them as churning at 24 months (also biased). Survival analysis treats the two types of observation differently in the likelihood.

For a churned customer at time $t_i$ , the likelihood contribution is the probability density $f(t_i)$ : we observed this exact event time. For a censored customer observed until time $c_i$ , the contribution is the survival probability $S(c_i) = P(T > c_i)$ : all we know is they lasted at least this long.

The total log-likelihood combines both pieces:

This is exactly how our PyMC model works. The pm.Gumbel line handles the first sum (uncensored density). The pm.Potential line handles the second sum (censored survival).

Why Gumbel? The Weibull-Gumbel Connection

The Weibull distribution is the workhorse of survival analysis because it models flexible hazard rates: increasing, decreasing, or constant over time. But working with the Weibull directly is numerically awkward for regression.

Here's the trick. If $T \sim \text{Weibull}(k, \lambda)$ , then $Y = \log T$ follows a Gumbel distribution:

where $\mu = \log \lambda$ is the location and $\sigma = 1/k$ is the scale. This is the accelerated failure time (AFT) formulation: covariates shift $\mu$ , effectively accelerating or decelerating time. We write the linear predictor as:

A positive $\alpha_1$ means higher spending shifts log-time to the right (longer survival). A negative $\alpha_2$ means more support tickets shift it left (faster churn). The coefficients have a direct interpretation: a one-unit increase in $x_j$ multiplies the median survival time by $\exp(\alpha_j)$ .

`pm.Potential`: Telling PyMC About Partial Information

In our hierarchical regression post, every observation contributed a full likelihood term through pm.Normal(..., observed=y). Censored observations are different: they don't have a fully observed outcome. They only contribute through the survival function.

pm.Potential('name', value) adds value directly to the model's log-posterior. For censored data, we pass the log-survival probability:

y_cens = pm.Potential('y_cens',
    gumbel_log_sf(log_observed_time[censored], eta[censored], s))

Think of it this way. For a churned customer, we say "we observed them leave at time $t$ " (standard likelihood). For an active customer, we say "all we know is they're still here after $c$ months" (survival function).

MCMC Diagnostics

Before trusting the results, verify the sampler converged:

az.plot_trace(trace, var_names=['alpha', 's'])
plt.tight_layout()

Check the same three diagnostics we covered in the hierarchical regression post: chains should look like "hairy caterpillars" (good mixing), R-hat below 1.01 (convergence), and effective sample size above 400 per chain (low autocorrelation).

Survival Curves from the Posterior

The payoff of a Bayesian AFT model is individual survival curves with uncertainty bands. For any customer profile, we compute the survival probability at each time point across all posterior samples:

t_grid = np.linspace(0.5, 36, 200)
log_t_grid = np.log(t_grid)

# Extract posterior samples
alpha_post = trace.posterior['alpha'].values.reshape(-1, 3)
s_post = trace.posterior['s'].values.flatten()

profiles = {
    'High-value (spend +1.5σ, tickets −1σ)': (1.5, -1.0, '#2196F3'),
    'Average customer':                        (0.0,  0.0, '#FF9800'),
    'At-risk (spend −1.5σ, tickets +2σ)':     (-1.5,  2.0, '#F44336'),
}

fig, ax = plt.subplots(figsize=(8, 5))
for label, (sp, tk, color) in profiles.items():
    eta_post = alpha_post[:, 0] + alpha_post[:, 1] * sp + alpha_post[:, 2] * tk
    survival = np.zeros((len(eta_post), len(t_grid)))
    for i in range(len(eta_post)):
        z = (log_t_grid - eta_post[i]) / s_post[i]
        survival[i] = 1 - np.exp(-np.exp(-z))
    mean_surv = survival.mean(axis=0)
    lower = np.percentile(survival, 3, axis=0)
    upper = np.percentile(survival, 97, axis=0)
    ax.plot(t_grid, mean_surv, color=color, lw=2, label=label)
    ax.fill_between(t_grid, lower, upper, color=color, alpha=0.15)

ax.set_xlabel('Months since signup')
ax.set_ylabel('Survival probability')
ax.set_title('Predicted Survival Curves by Customer Profile')
ax.legend(loc='upper right', fontsize=9)
ax.set_xlim(0, 36)
ax.set_ylim(0, 1.05)
plt.tight_layout()

Each curve shows the model's predicted probability that a customer with those characteristics survives beyond a given time. The high-value customer has a much flatter curve: their predicted median lifetime exceeds 36 months. The at-risk customer (low spend, many support tickets) has a steep drop-off with a median around 5 months.

Notice the uncertainty bands widen at longer times, especially for the at-risk profile. Fewer customers with those characteristics survive that long, so the model has less data to constrain the prediction.

Going Deeper

Covariates in the Scale Too

The model above uses a constant scale parameter $s$ for all customers. The original code I adapted goes further by making the scale covariate-dependent:

This means the shape of the Weibull hazard varies across customers. A customer might have both a longer expected lifetime (larger $\eta$ ) and more predictable survival (smaller $s$ ). In PyMC:

with pm.Model() as weibull_aft_hetero:
    # Location coefficients
    alpha = pm.Normal('alpha', mu=0, sigma=2, shape=3)
    # Scale coefficients (matching original code's rho priors)
    rho = pm.Normal('rho', mu=0, sigma=2, shape=3)

    eta = alpha[0] + alpha[1] * spend_std + alpha[2] * tickets_std
    s = pm.math.exp(rho[0] + rho[1] * spend_std + rho[2] * tickets_std)

    y_obs = pm.Gumbel('y_obs', mu=eta[~censored], beta=s[~censored],
                       observed=log_observed_time[~censored])
    y_cens = pm.Potential('y_cens',
        gumbel_log_sf(log_observed_time[censored], eta[censored], s[censored]))

    trace_hetero = pm.sample(1000, tune=2000, cores=4, chains=4,
                             random_seed=42, target_accept=0.9)

This is faithful to the aft_model_factory_explicit function in the original code, which uses separate rho_interc, rho_coeff1, rho_coeff2 parameters for the Gumbel scale. The exp link ensures $s_i > 0$ for every customer.

Weibull vs Log-Logistic: Which Tail Shape?

The Weibull model assumes the hazard rate is monotonic: always increasing, always decreasing, or constant. But some churn patterns are non-monotonic. New users might have high churn risk initially (they haven't found value yet), which drops as they engage, then rises again as they outgrow the product.

The Log-Logistic AFT model handles this. In log-time, the Log-Logistic corresponds to a Logistic distribution, just as the Weibull corresponds to a Gumbel. The swap is straightforward:

def logistic_log_sf(y, mu, sigma):
    """Log survival function of the Logistic distribution."""
    return -pt.softplus((y - mu) / sigma)

with pm.Model() as loglogistic_aft:
    alpha = pm.Normal('alpha', mu=0, sigma=2, shape=3)
    log_s = pm.Normal('log_s', mu=0, sigma=1)
    s = pm.Deterministic('s', pm.math.exp(log_s))

    eta = alpha[0] + alpha[1] * spend_std + alpha[2] * tickets_std

    y_obs = pm.Logistic('y_obs', mu=eta[~censored], s=s,
                         observed=log_observed_time[~censored])
    y_cens = pm.Potential('y_cens',
        logistic_log_sf(log_observed_time[censored], eta[censored], s))

    trace_ll = pm.sample(1000, tune=2000, cores=4, chains=4,
                         random_seed=42, target_accept=0.9)

Compare the two models using LOO-CV (leave-one-out cross-validation) with ArviZ:

weibull_loo = az.loo(trace)
ll_loo = az.loo(trace_ll)
print(az.compare({'Weibull': trace, 'Log-Logistic': trace_ll}))

Since our synthetic data was generated from a Weibull distribution, the Weibull model should win. On real data, the comparison often reveals which tail shape better captures your customers' churn dynamics.

The Cox Proportional Hazards Alternative

Survival analysis has a dominant semi-parametric approach: the Cox proportional hazards (PH) model. It doesn't assume a distribution for the baseline hazard, only that covariates multiply the hazard by a constant factor. This flexibility made it ubiquitous in clinical trials.

So why choose a parametric Bayesian AFT model? Three reasons:

Full predictive distributions. The Cox model gives hazard ratios, but producing survival curves requires additional estimation of the baseline hazard. Our Bayesian AFT model gives survival curves with uncertainty bands directly from the posterior.
Small samples and heavy censoring. With many active customers, the Cox model's partial likelihood can be imprecise. Bayesian priors stabilise estimates, especially for rare covariates. This is the same principle of "borrowing strength" we explored in the hierarchical regression post.
Natural extension. PyMC models compose freely. Adding group structure (churn by subscription tier), time-varying covariates, or custom likelihoods is straightforward. The next post in this series demonstrates exactly this with a one-inflated Beta regression.

When NOT to Use Bayesian AFT

If the proportional hazards assumption holds and your dataset is large (tens of thousands of events), the Cox model is faster and assumption-lighter. If you have time-varying covariates that change during a customer's lifetime (e.g., monthly usage patterns), the standard AFT formulation doesn't handle them naturally; you'd need a piecewise approach or a joint model.

Computational cost matters too. Our 1,000-customer model samples in a few minutes, but production datasets with millions of rows would require approximations like variational inference or mini-batch MCMC.

Where This Comes From

Cox (1972): Proportional Hazards

The modern era of survival analysis began with David Cox's 1972 paper "Regression Models and Life-Tables." Cox introduced the proportional hazards model:

where $h_0(t)$ is an unspecified baseline hazard. The genius was leaving $h_0$ unspecified and estimating $\boldsymbol{\beta}$ through the partial likelihood, which depends only on the order of events, not their exact times. This paper has been cited over 65,000 times and remains the most-used method in clinical trials.

"The important practical point is that [the partial likelihood] does not require specification of $h_0(t)$ ." (Cox, 1972)

Our AFT model takes a different path: we specify a distribution (Weibull or Log-Logistic), which enables direct time predictions. This parametric assumption is both a strength (more powerful inference when correct) and a weakness (biased inference when wrong).

Buckley and James (1979): Accelerated Failure Time

The AFT framework was formalised by Miles Buckley and Ian James in 1979. Their key insight was that the AFT model has a direct linear regression interpretation:

where $\epsilon_i$ follows a known distribution (Gumbel for Weibull, Logistic for Log-Logistic). The coefficients $\alpha_j$ have a clean meaning: a one-unit increase in $x_j$ multiplies the median survival time by $\exp(\alpha_j)$ . This is why it's called "accelerated failure time": covariates speed up or slow down the passage of time.

Wei (1992): AFT as an Alternative

L. J. Wei's 1992 paper "The Accelerated Failure Time Model: A Useful Alternative to the Cox Regression Model in Survival Analysis" made the case for AFT models as a practical complement to Cox PH. Wei showed that AFT models are more robust to omitted covariates and provide more interpretable effect sizes.

"When the acceleration factor is constant over time, the AFT model provides a simple and clinically meaningful summary of the survival experience." (Wei, 1992)

Handling Censoring in PyMC

The pm.Potential approach for censored data follows directly from the likelihood factorisation. For a dataset with observed and censored outcomes:

Taking logs, the uncensored terms give the standard log-likelihood (handled by pm.Gumbel or pm.Logistic). The censored terms give log-survival values (handled by pm.Potential). This pattern appears throughout the PyMC survival analysis examples and extends naturally to interval censoring and left censoring by swapping the survival function for the appropriate probability term.

Interactive Tools

Kaplan-Meier Calculator — Estimate survival curves and compare groups interactively
Medical Statistics Calculator — Compute sensitivity, specificity, and other diagnostic metrics

Hierarchical Bayesian Regression with PyMC: The first post in this PyMC series, covering partial pooling and MCMC diagnostics with ArviZ.
MCMC Island Hopping: Understanding Metropolis-Hastings: How the NUTS sampler that powers PyMC explores high-dimensional posteriors.
From MLE to Bayesian Inference: The conceptual foundation for priors, posteriors, and why Bayesian estimates outperform point estimates.

Frequently Asked Questions

What is right-censoring and why does it matter?

Right-censoring occurs when you know a subject survived at least until a certain time, but not the actual event time. In churn analysis, active customers are right-censored because they have not yet churned. Ignoring them biases your model toward shorter lifetimes, since you only learn from customers who already left. Survival analysis handles censoring properly by using the survival function for these partial observations.

What is the difference between the Cox model and an AFT model?

The Cox proportional hazards model is semi-parametric: it leaves the baseline hazard unspecified and estimates how covariates multiply the hazard rate. The accelerated failure time (AFT) model is fully parametric: it assumes a specific distribution (such as Weibull) and models how covariates accelerate or decelerate time to event. AFT coefficients have a direct interpretation as multipliers on median survival time, while Cox coefficients are hazard ratios.

What does pm.Potential do in PyMC?

pm.Potential adds an arbitrary log-probability term directly to the model's log-posterior. For censored observations, there is no fully observed outcome to pass to a standard likelihood. Instead, you compute the log-survival probability and add it via pm.Potential, telling PyMC that these customers survived at least this long without specifying when they will actually churn.

How do I choose between Weibull and Log-Logistic distributions?

Use Weibull when you expect the hazard rate to be monotonic, either always increasing, always decreasing, or constant over time. Use Log-Logistic when the hazard may be non-monotonic, such as high initial churn that drops as users engage and then rises again later. You can compare the two formally using LOO-CV (leave-one-out cross-validation) in ArviZ.

How many customers do I need for a Bayesian survival model?

Bayesian models can work with surprisingly small datasets because priors regularise the estimates, but a practical minimum is a few hundred observations with at least 50 to 100 uncensored events. With heavy censoring (over 80% still active), the model has less information about event times, so you may need a larger sample or more informative priors to get precise estimates.

Can I add time-varying covariates to a Bayesian AFT model?

The standard AFT formulation assumes covariates are fixed at baseline and does not naturally handle features that change during a customer's lifetime, such as monthly usage patterns. For time-varying covariates, you would need a piecewise AFT approach that splits each customer's timeline into intervals, or a joint model that links the longitudinal covariate process with the survival outcome.

Hierarchical Bayesian Regression with PyMC: When Groups Share Strength

Berkan Sesen — Sun, 26 Apr 2026 12:43:53 +0000

A multi-line insurer writes auto, home, commercial property, and a dozen other policy types under one roof. Some lines see thousands of claims a year; others might see 50. Every actuary faces the same dilemma: train a separate pricing model for each line and the small ones are pure noise, or pool everything together and pretend a warehouse fire looks like a fender bender. Either way, you lose.

Hierarchical Bayesian regression offers a third way. Each group gets its own parameters, but those parameters are drawn from a shared population distribution. Groups with plenty of data stay close to their own estimates. Groups with little data get "pulled" toward the population average, borrowing statistical strength from the larger groups. This effect is called shrinkage, and it's one of the most elegant ideas in statistics.

By the end of this post, you'll build a hierarchical Bayesian regression model in PyMC, compare it against pooled and unpooled alternatives, and see shrinkage in action on synthetic insurance data.

Let's Build It

First, let's see the hierarchical model in action. Click the badge below to open the full interactive notebook:

We'll generate synthetic insurance claim data for three policy types with deliberately unbalanced sample sizes, then fit a hierarchical model.

import numpy as np
import pandas as pd
import pymc as pm
import arviz as az
import matplotlib.pyplot as plt

np.random.seed(42)

# Three policy types: lots of Auto data, moderate Home, very little Commercial
groups = {
    'Auto':       {'n': 500, 'intercept': 7.5, 'slope': 0.30},
    'Home':       {'n': 300, 'intercept': 8.2, 'slope': 0.50},
    'Commercial': {'n':  50, 'intercept': 9.0, 'slope': 0.70},
}

records = []
for i, (name, p) in enumerate(groups.items()):
    x = np.random.normal(12, 1.5, p['n'])  # log property value (~$160k median)
    noise = np.random.normal(0, 0.8, p['n'])
    y = p['intercept'] + p['slope'] * (x - 12) + noise
    for j in range(p['n']):
        records.append({
            'policy_type': name, 'group_idx': i,
            'log_property_value': x[j], 'log_claim_severity': y[j],
        })

df = pd.DataFrame(records)

Each policy type has a different intercept and slope, but Commercial has just 50 data points. Now let's fit the hierarchical model:

n_types = len(groups)
idx = df['group_idx'].values
x_centered = df['log_property_value'].values - 12  # center the predictor

with pm.Model() as hierarchical_model:
    # Hyperpriors: the "population" distribution that groups are drawn from
    mu_alpha = pm.Normal('mu_alpha', mu=8, sigma=2)
    sigma_alpha = pm.HalfNormal('sigma_alpha', sigma=2)
    mu_beta = pm.Normal('mu_beta', mu=0, sigma=2)
    sigma_beta = pm.HalfNormal('sigma_beta', sigma=2)

    # Group-level parameters, drawn from the population
    alpha = pm.Normal('alpha', mu=mu_alpha, sigma=sigma_alpha, shape=n_types)
    beta = pm.Normal('beta', mu=mu_beta, sigma=sigma_beta, shape=n_types)

    # Observation noise
    sigma = pm.HalfNormal('sigma', sigma=2)

    # Linear model
    mu = alpha[idx] + beta[idx] * x_centered
    y_obs = pm.Normal('y_obs', mu=mu, sigma=sigma,
                      observed=df['log_claim_severity'].values)

    # Sample the posterior
    hierarchical_trace = pm.sample(
        1000, tune=1500, cores=4, chains=4,
        random_seed=42, target_accept=0.9,
    )

# Summarise the results
print(az.summary(hierarchical_trace, var_names=['alpha', 'beta', 'sigma']))

You just estimated three group-specific regression lines (one per policy type) while letting them share statistical strength through a common population distribution. The Commercial group, despite having only 50 claims, gets a stable estimate because it borrows information from Auto and Home.

What Just Happened?

The Three Pooling Strategies

To understand why the hierarchical model is special, let's compare it against the two extreme alternatives.

Complete pooling ignores group differences entirely. One intercept, one slope for all 850 data points:

with pm.Model() as pooled_model:
    alpha = pm.Normal('alpha', mu=8, sigma=5)
    beta = pm.Normal('beta', mu=0, sigma=5)
    sigma = pm.HalfNormal('sigma', sigma=2)

    mu = alpha + beta * x_centered
    y_obs = pm.Normal('y_obs', mu=mu, sigma=sigma,
                      observed=df['log_claim_severity'].values)

    pooled_trace = pm.sample(1000, tune=1000, cores=4, chains=4, random_seed=42)

No pooling treats each group as completely independent. Three separate intercepts, three separate slopes, with no shared information:

with pm.Model() as unpooled_model:
    alpha = pm.Normal('alpha', mu=8, sigma=5, shape=n_types)
    beta = pm.Normal('beta', mu=0, sigma=5, shape=n_types)
    sigma = pm.HalfNormal('sigma', sigma=2)

    mu = alpha[idx] + beta[idx] * x_centered
    y_obs = pm.Normal('y_obs', mu=mu, sigma=sigma,
                      observed=df['log_claim_severity'].values)

    unpooled_trace = pm.sample(1000, tune=1000, cores=4, chains=4, random_seed=42)

The comparison reveals the key insight. Complete pooling gives a single line dominated by Auto and Home (which together make up 94% of the data), systematically underestimating Commercial's higher intercept and steeper slope. No pooling gives each group its own line, but Commercial's estimate is noisy because it only has 50 points. Partial pooling (the hierarchical model) sits between the two: each group gets its own line, but the lines are gently pulled toward the population average. Groups with little data get pulled more.

How Hyperpriors Create Partial Pooling

The magic ingredient is the hyperpriors: mu_alpha, sigma_alpha, mu_beta, sigma_beta. These define a "population distribution" from which group-level parameters are drawn.

Think of $\mu_\alpha$ as the average intercept across all policy types, and $\sigma_\alpha$ as how much the types are allowed to differ. If the data supports large differences, $\sigma_\alpha$ will be large and each group behaves almost independently (like no pooling). If the groups are similar, $\sigma_\alpha$ shrinks and the group estimates collapse toward the population mean (like complete pooling).

The sampler learns $\sigma_\alpha$ from the data itself. You don't have to choose between pooling and no pooling; the model figures out the right amount of sharing automatically.

Shrinkage: The Key Insight

Shrinkage is the defining feature of hierarchical models. Compare each group's raw sample mean (what you'd get from no pooling) to its hierarchical posterior mean:

Commercial's intercept gets pulled the most toward the population mean, because it has the least data and therefore the most uncertainty. Auto barely moves, because 500 data points leave little room for the prior to override the evidence. This is exactly the Bayesian compromise between prior and data that we explored in From MLE to Bayesian Inference.

MCMC Diagnostics with ArviZ

Before trusting the results, we need to verify the sampler converged. ArviZ provides the standard toolkit:

az.plot_trace(hierarchical_trace, var_names=['alpha', 'beta', 'sigma'])
plt.tight_layout()

Three things to check:

Trace mixing: The chains should look like "hairy caterpillars", bouncing randomly around a stable mean. If a chain gets stuck or drifts, something is wrong.
R-hat (the Gelman-Rubin statistic): Should be below 1.01 for every parameter. Values above 1.1 indicate the chains haven't converged to the same distribution.
Effective sample size (ESS): Should be at least 400 per chain. Low ESS means the samples are highly autocorrelated and the posterior estimates are unreliable.

summary = az.summary(hierarchical_trace, var_names=['alpha', 'beta', 'sigma'])
print(summary[['mean', 'sd', 'hdi_3%', 'hdi_97%', 'r_hat', 'ess_bulk']])

If you've worked through our MCMC Metropolis-Hastings tutorial, you'll recognise the core idea: the sampler explores the posterior by proposing moves and accepting or rejecting them. PyMC uses the NUTS sampler (No U-Turn Sampler), a sophisticated variant of Hamiltonian Monte Carlo that automatically tunes step sizes and trajectory lengths.

Going Deeper

Why Not a Normal Likelihood?

The model above uses a Normal likelihood, which assumes claim amounts are symmetric around the mean. In practice, insurance claims are heavy-tailed: most claims are small, but a few are enormous. The original code I adapted for this tutorial used a Laplace likelihood to handle this:

The Laplace distribution has heavier tails than the Normal and is more robust to outliers. In PyMC, swapping the likelihood is a single line change:

# Replace:  pm.Normal('y_obs', mu=mu, sigma=sigma, observed=y)
# With:     pm.Laplace('y_obs', mu=mu, b=b, observed=y)

Modelling the Spread Too: Heteroscedastic Regression

The original code goes further. It models both the location $\mu$ and the scale $b$ of the Laplace distribution as functions of the covariates. This is heteroscedastic regression: the amount of noise varies across observations.

The $\exp$ ensures both $\mu$ and $b$ are positive (claim severity can't be negative). Each $\beta$ and $\gamma$ coefficient gets its own hierarchical structure:

with pm.Model() as full_model:
    n_groups = 3  # policy types

    # Hyperpriors for intercept
    beta0_mu = pm.Normal('beta0_mu', mu=0, sigma=1)
    beta0_sig = pm.InverseGamma('beta0_sig', alpha=2, beta=5)

    # Group-level intercepts
    beta0 = pm.Normal('beta0', mu=beta0_mu, sigma=beta0_sig, shape=n_groups)

    # ... repeat for each coefficient and for gamma (scale) parameters ...

    mu = pm.math.exp(beta0[group_idx] + beta1[group_idx] * pm.math.log(X[:, 0]) + ...)
    b  = pm.math.exp(gamma0[group_idx] + gamma1[group_idx] * pm.math.log(X[:, 0]) + ...)

    y_obs = pm.Laplace('y_obs', mu=mu, b=b, observed=y)

Notice the pm.InverseGamma hyperprior for the variance parameters. The InverseGamma is the conjugate prior for Normal variance, making it a natural choice. With alpha=2, beta=5, it places mass on moderate variance values while allowing large ones.

The Three-Tier Model

The code also contains a three-tier hierarchy. Instead of just grouping by policy type, it nests policy type within region:

Population → Policy Type → (Region × Policy Type)

At the top level, hyper-hyperpriors define the global population. At the middle level, each policy type gets its own parameters drawn from the population. At the bottom level, each (region, policy type) combination gets parameters drawn from its policy type's distribution. The group-level parameters become 2D arrays with shape (n_regions, n_policy_types):

# Hyper-hyperpriors (population level)
beta0_mu_mu = pm.Normal('beta0_mu_mu', mu=0, sigma=1)
beta0_mu_sig = pm.InverseGamma('beta0_mu_sig', alpha=2, beta=5)

# Hyperpriors (policy type level)
beta0_mu = pm.Normal('beta0_mu', mu=beta0_mu_mu, sigma=beta0_mu_sig, shape=n_types)
beta0_sig = pm.InverseGamma('beta0_sig', alpha=2, beta=5, shape=n_types)

# Priors (region × policy type level)
beta0 = pm.Normal('beta0', mu=beta0_mu, sigma=beta0_sig, shape=(n_regions, n_types))

This allows a commercial policy in an urban area to differ from one in a suburban area, while both borrow strength from the overall commercial distribution, which itself borrows from the global population.

When Not to Use Hierarchical Models

Hierarchical models aren't always necessary. If every group has plenty of data (thousands of observations), no pooling gives nearly identical results to partial pooling because the data overwhelms the prior. The hierarchy adds complexity and sampling time for little benefit.

They can also struggle with very few groups. With only 2 groups, the hyperprior variance $\sigma_\alpha$ is estimated from just 2 data points (the two group-level parameters), making it unreliable. Most practitioners suggest hierarchical models shine with 5 or more groups, though the exact threshold depends on within-group sample sizes.

Where This Comes From

Lindley and Smith (1972)

The mathematical foundation was laid by Dennis Lindley and Adrian Smith in their 1972 paper "Bayes Estimates for the Linear Model." They formalised the multi-level Normal model:

The key result: the posterior mean of $\boldsymbol{\theta}$ is a matrix-weighted average of the group-specific MLE and the prior mean. Groups with more data (higher precision in $C^{-1}$ ) weight their own MLE more heavily; groups with less data lean more on the prior. This is the formal statement of shrinkage.

Efron and Morris (1977): The James-Stein Connection

The frequentist justification for shrinkage came from an unexpected direction. In 1977, Brad Efron and Carl Morris showed that the James-Stein estimator (which shrinks group means toward the grand mean) dominates the usual sample means in terms of total squared error, for three or more groups simultaneously. This was a shocking result: even if the groups have nothing in common, shrinking toward their average reduces total estimation error.

"The James-Stein estimator achieves a smaller total mean squared error than the individual sample means, for any configuration of the true means, provided there are three or more groups."

The hierarchical Bayesian model produces estimates that are closely related to the James-Stein estimator. The Bayesian framework provides a natural explanation: when data is scarce, it's rational to hedge toward the population average rather than fully committing to a noisy local estimate.

Gelman and Hill (2006)

The practical handbook for hierarchical models is Andrew Gelman and Jennifer Hill's Data Analysis Using Regression and Multilevel/Hierarchical Models. Chapter 12 presents the exact three-model comparison we built above (complete pooling, no pooling, partial pooling) using radon measurements across US counties. Their formulation uses the non-centred parameterisation:

This reparameterisation often improves MCMC sampling efficiency because the sampler explores a standard Normal geometry rather than a funnel-shaped one. PyMC can apply this transformation automatically, but it's worth knowing when your model has divergences.

Gelman et al.'s Bayesian Data Analysis (3rd edition, 2013) provides the full mathematical treatment in Chapter 5, including the relationship between hierarchical Bayes, empirical Bayes, and the James-Stein estimator.

Interactive Tools

Bayes' Theorem Calculator — Explore Bayesian updating interactively before diving into hierarchical models
A/B Test Calculator — See Bayesian hypothesis testing in action, a common application of hierarchical models

From MLE to Bayesian Inference: The conceptual foundation for priors, posteriors, and why Bayesian estimates beat point estimates.
MCMC Island Hopping: Understanding Metropolis-Hastings: How the sampler that powers PyMC actually explores the posterior distribution.
Linear Regression: Five Ways: The non-hierarchical regression baseline that this post extends with group structure.

Frequently Asked Questions

What is hierarchical Bayesian regression?

Hierarchical (or multilevel) regression models data that is naturally grouped (students within schools, patients within hospitals) by allowing parameters to vary across groups while sharing a common prior distribution. This "partial pooling" approach borrows strength across groups, producing better estimates for small groups than fitting each group independently.

What is the difference between complete pooling, no pooling, and partial pooling?

Complete pooling ignores group differences entirely (one model for all). No pooling fits a separate model per group (no information sharing). Partial pooling (hierarchical) sits in between: each group gets its own parameters, but they are pulled towards a shared distribution. This is especially valuable when some groups have very few observations.

Why use PyMC for hierarchical models?

PyMC uses MCMC sampling to handle the complex posterior distributions that hierarchical models produce. It naturally propagates uncertainty through all levels of the hierarchy. Frequentist alternatives (like lme4 in R) can fit similar models but do not provide the same rich uncertainty quantification or flexibility for custom model structures.

How do I diagnose convergence in PyMC?

Check the trace plots for good mixing (no trends, no stuck chains), verify that R-hat values are close to 1.0 (below 1.01), and ensure effective sample sizes are sufficiently large (at least 400 per chain). Divergent transitions indicate the sampler is struggling with the posterior geometry and may require reparameterisation.

When should I use a hierarchical model instead of a standard regression?

Use hierarchical models whenever your data has a natural grouping structure and you want to make inferences about individual groups. They are especially valuable when group sizes are unequal: small groups benefit from borrowing strength, and large groups are barely affected by the pooling. If all groups have abundant data, the results will be similar to fitting separate models.

Solving CartPole Without Gradients: Simulated Annealing

Berkan Sesen — Thu, 23 Apr 2026 07:51:02 +0000

In the previous post, we solved CartPole using the Cross-Entropy Method: sample 200 candidate policies, keep the best 40, refit a Gaussian, repeat. It worked beautifully, reaching a perfect score of 500 in 50 iterations. But 200 candidates per iteration means 10,000 total episode evaluations. That got me wondering: do we really need a population of 200 to find four good numbers?

The original code that inspired this post took a radically simpler approach. Instead of maintaining a population, it kept a single set of parameters and perturbed them once per iteration. If the perturbation improved the score, it was accepted and the perturbation range was shrunk. That's it. No population, no distribution fitting, no gradients. The comment in the source file read: "its like simulated annealing." By the end of this post, you'll implement this algorithm from scratch, solve CartPole-v1 with a perfect 500 score, and understand how it connects to the rich theory of simulated annealing.

Let's Build It

Click the badge to open the interactive notebook:

Here's the complete implementation. Like CEM, we use a linear policy with 4 parameters (one per observation dimension). But instead of sampling a population, we perturb a single solution:

import numpy as np
import gymnasium as gym

def evaluate_policy(env_name, theta, n_episodes=10):
    """Run multiple episodes with a linear policy and return the average reward."""
    total_reward = 0
    for _ in range(n_episodes):
        env = gym.make(env_name)
        obs, _ = env.reset()
        episode_reward = 0
        done = False
        while not done:
            action = 1 if np.dot(theta, obs) > 0 else 0
            obs, reward, terminated, truncated, _ = env.step(action)
            episode_reward += reward
            done = terminated or truncated
        env.close()
        total_reward += episode_reward
    return total_reward / n_episodes

def simulated_annealing(env_name, n_params, n_iter=80, n_eval_episodes=10,
                        alpha=1.0, decay=0.9):
    """Hill climbing with annealing step size for policy search."""
    best_theta = np.zeros(n_params)
    best_score = evaluate_policy(env_name, best_theta, n_eval_episodes)

    for i in range(n_iter):
        # Perturb current best (uniform noise scaled by alpha)
        perturbation = (np.random.rand(n_params) - 0.5) * alpha
        candidate = best_theta + perturbation

        # Evaluate candidate over multiple episodes
        score = evaluate_policy(env_name, candidate, n_eval_episodes)

        # Accept only if better, then shrink step size
        if score > best_score:
            best_theta = candidate
            best_score = score
            alpha *= decay

        print(f"Iter {i+1:3d} | Score: {score:.1f} | Best: {best_score:.1f} | Alpha: {alpha:.4f}")

    return best_theta

best_theta = simulated_annealing('CartPole-v1', n_params=4)
# Iter   1 | Score:   9.6 | Best:   9.6 | Alpha: 1.0000
# Iter   9 | Score: 128.7 | Best: 128.7 | Alpha: 0.6561
# Iter  14 | Score: 314.2 | Best: 314.2 | Alpha: 0.5314
# Iter  24 | Score: 465.7 | Best: 465.7 | Alpha: 0.4783
# Iter  41 | Score: 500.0 | Best: 500.0 | Alpha: 0.3874

Perfect score in 41 iterations. Let's verify with 100 evaluation episodes:

scores = [evaluate_policy('CartPole-v1', best_theta, n_episodes=1) for _ in range(100)]
print(f"Mean: {np.mean(scores):.0f} +/- {np.std(scores):.0f}")
# Mean: 496 +/- 12

Four parameters, zero gradients, 800 total episode evaluations. Compare that to CEM's 10,000 episodes or REINFORCE's 5,000.

What Just Happened?

The algorithm maintains a single candidate solution and improves it through a cycle of perturb, evaluate, and accept. Here's the full loop:

Let's walk through each piece.

The Linear Policy

Just like in the CEM post, CartPole has a 4-dimensional observation vector (cart position, cart velocity, pole angle, pole angular velocity). Our policy is a simple dot product:

action = 1 if np.dot(theta, obs) > 0 else 0

This is a linear classifier: push right if the weighted sum of observations is positive, push left otherwise. The entire "intelligence" of the agent lives in four numbers.

Multi-Episode Evaluation

The original code's key insight (noted in a comment: "key thing was to figure out that you need to do 10 tests per point") is to evaluate each candidate over 10 episodes and average the scores. CartPole has stochastic initial conditions, so a single episode can be misleading. A policy might score 500 on one lucky initialisation and 50 on the next. Averaging over 10 episodes gives a stable estimate of true quality.

score = evaluate_policy(env_name, candidate, n_eval_episodes=10)

The Perturbation Step

Each iteration, we perturb the current best parameters with uniform noise scaled by alpha:

perturbation = (np.random.rand(n_params) - 0.5) * alpha
candidate = best_theta + perturbation

When alpha=1.0, each parameter can change by up to $\pm 0.5$ . As alpha shrinks, the perturbations get smaller, focusing the search around the current best.

Accept and Anneal

Here's the crucial part. We only accept improvements, and we only shrink the step size when we find one:

if score > best_score:
    best_theta = candidate
    best_score = score
    alpha *= decay  # Shrink step size by 10%

This is an adaptive cooling schedule. If the algorithm keeps finding improvements, alpha decays quickly ( $0.9^9 \approx 0.39$ after 9 improvements). If it gets stuck, alpha stays large, maintaining exploration. The algorithm found 9 improvements out of 80 iterations, ending with $\alpha = 0.387$ .

The Training Curve

The staircase pattern tells the story. Each vertical jump is an accepted improvement; each flat region is the algorithm searching without finding anything better:

fig, ax1 = plt.subplots(figsize=(8, 4))
ax1.scatter(iters[1:], candidate_scores[1:],
            c=['#2ecc71' if a else '#e74c3c' for a in accepted[1:]],
            alpha=0.5, s=20, label='Candidates')
ax1.plot(iters, best_scores, 'b-', linewidth=2, label='Best score')
ax1.axhline(y=500, color='k', linestyle=':', alpha=0.3, label='Max possible (500)')

ax2 = ax1.twinx()
ax2.plot(iters, alphas, 'k--', alpha=0.4, linewidth=1, label='Step size (α)')
ax2.set_ylabel('Step size (α)', color='gray')

Green dots are accepted candidates (improvements); red dots are rejected ones. The dashed grey line shows the step size $\alpha$ shrinking on the secondary axis. Notice how the red dots cluster higher as the search progresses, because even rejected perturbations from a good solution tend to produce decent policies.

Going Deeper

Hill Climbing vs True Simulated Annealing

Let's be precise about what our algorithm is. The original code's comment called it "like simulated annealing," and that's accurate, but with an important distinction.

Our algorithm (hill climbing with annealing step size):

Accepts only improvements
Shrinks the step size when an improvement is found
Never accepts a worse solution

True simulated annealing:

Accepts improvements always
Accepts worse solutions with probability $e^{-\Delta E / T}$
Shrinks the temperature $T$ on a fixed schedule

The difference is in how they handle worse solutions. True SA occasionally accepts a downhill move, which allows it to escape local optima. Our algorithm never does, which makes it a strict hill climber. The "annealing" part is only in the step size, not in the acceptance criterion.

For CartPole with a 4-parameter linear policy, this distinction doesn't matter: the reward landscape is smooth enough that hill climbing works. For harder problems with many local optima, true SA's ability to escape traps becomes essential.

If you've read the MCMC Metropolis-Hastings post, the acceptance criterion should look familiar. The Metropolis acceptance probability $\min(1, e^{-\Delta E / T})$ is exactly what true SA uses. In MCMC, we want to sample from a distribution; in SA, we want to find its peak. Same mechanism, different goal.

The Cooling Schedule

Our algorithm uses a multiplicative decay: $\alpha_{t+1} = 0.9 \cdot \alpha_t$ on each improvement. This creates a geometric sequence:

where $k$ is the number of improvements found. After 9 improvements, $\alpha = 0.9^9 \approx 0.387$ .

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
ax1.plot(iters, alphas, 'b-', linewidth=2)       # Alpha vs iterations
ax2.plot(steps, 0.9**steps, linewidth=2)          # Geometric decay curves

The left panel shows alpha over iterations, with green bands marking accepted improvements. The right panel compares different decay rates. A faster decay ( $\gamma = 0.8$ ) converges to fine-tuning quickly but risks getting stuck. A slower decay ( $\gamma = 0.95$ ) explores longer but takes more iterations to refine. The original code's choice of 0.9 strikes a reasonable balance.

What makes our schedule adaptive is that it only decays on improvement. Traditional SA uses fixed schedules (logarithmic, linear, or exponential decay in wall-clock time). Our variant keeps $\alpha$ large during plateaus, naturally spending more time exploring when stuck and more time refining when making progress.

SA vs CEM: One Climber vs a Search Party

The Cross-Entropy Method we built last time and simulated annealing sit at opposite ends of the derivative-free spectrum:

Aspect	Simulated Annealing	CEM
Search strategy	Single point, local perturbations	Population of 200, distribution fitting
Episodes per iteration	10	200 (200 candidates x 1 each)
Total episodes to solve CartPole	~800	~10,000 (200 x 50 iterations)
Information used	"Is this better than the best?" (1 bit)	Full reward ranking of all candidates
Robustness	Seed-dependent; some runs may fail	Highly robust; population averages out noise
Parallelisable	No (sequential by nature)	Yes (all 200 evaluations are independent)

SA is like a single hiker exploring a mountain range, taking one step at a time and only moving to higher ground. CEM is like sending 200 hikers, ranking them by altitude, and teleporting the next batch to the region where the best ones clustered.

SA wins on sample efficiency (fewer total episodes) but loses on reliability. Run SA with a different random seed and you might need 20 iterations or 200. CEM's population averaging makes it much more consistent.

SA vs Random Search

How much does the "annealing" (building on previous improvements) actually help, compared to just sampling random policies each time?

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(iters, sa_best_scores, 'b-', linewidth=2, label='Simulated annealing')
ax.plot(iters, random_best_scores, 'r--', linewidth=2, label='Random search')

Random search samples a fresh random policy each iteration (uniform in $[-1, 1]^4$ ) and tracks the best one found. After 80 iterations, its best score is 387 vs SA's 500. Random search got lucky once (iteration 2) and found a decent policy early, but it can never refine it. SA's ability to make small improvements to an already-good solution is what pushes it from "decent" to "perfect."

Hyperparameters

Parameter	Value	Effect
`alpha`	1.0	Initial step size. Perturbations range in $[-0.5, 0.5]$ per parameter
`decay`	0.9	Step size multiplier on improvement. Lower = faster convergence, less exploration
`n_iter`	80	Total iterations. Our run converged at iteration 41
`n_eval_episodes`	10	Episodes per evaluation. More = less noise, more compute

The most sensitive parameter is decay. At 0.9, alpha halves after about 7 improvements. At 0.8, it halves after 4. Too aggressive and the step size collapses before finding a good solution; too conservative and you waste iterations on large perturbations when you're already close.

When NOT to Use This Approach

High-dimensional parameter spaces. A single perturbation in 1000 dimensions is unlikely to improve on the current best by chance. Population methods like CEM or genetic algorithms scale better
Multi-modal reward landscapes. Our hill climber can only find the nearest peak. If the global optimum is separated by a valley, you'll never reach it without true SA's downhill acceptance
When you need guarantees. SA is a heuristic. Even true SA only guarantees convergence to the global optimum with logarithmic cooling, which is impractically slow
When wall-clock time matters more than sample efficiency. SA is inherently sequential. CEM's 200 evaluations per iteration can run in parallel, making it faster on multi-core hardware despite using 12x more episodes

Where This Comes From

Simulated annealing was introduced independently by Scott Kirkpatrick, Daniel Gelatt, and Mario Vecchi at IBM Research in their 1983 Science paper "Optimization by Simulated Annealing", and by Vlasta Cerny in 1985. The name comes from the metallurgical process of annealing: heating a metal and then slowly cooling it to reduce defects in its crystal structure.

The Metallurgy Analogy

When you heat metal, atoms vibrate wildly and can escape local energy minima. As the temperature drops, atoms settle into increasingly stable configurations. If you cool slowly enough, the metal reaches its lowest-energy crystal state (the global optimum). Cool too fast and you get a brittle, disordered structure (a local optimum).

Kirkpatrick and colleagues mapped this physical process to combinatorial optimisation:

Metal atoms become candidate solutions
Energy becomes the cost function
Temperature becomes a control parameter that governs randomness

The Metropolis Connection

The acceptance criterion in true SA comes directly from the Metropolis algorithm (Metropolis, Rosenbluth, Rosenbluth, Teller, and Teller, 1953), originally designed for simulating atomic systems in statistical mechanics. At temperature $T$ , a new state with energy $E'$ is accepted from a current state with energy $E$ with probability:

At high $T$ , the exponential is close to 1, so almost any move is accepted (random exploration). At low $T$ , only improvements or tiny degradations are accepted (local refinement). As $T \to 0$ , the algorithm becomes pure hill climbing.

This is the same acceptance probability we explored in the Metropolis-Hastings post for MCMC sampling. The only difference: in MCMC, we maintain a high temperature to sample broadly; in SA, we lower it to converge on a peak. Same mechanism, different goals.

Our Variant vs Classical SA

Our implementation simplifies classical SA in two ways:

No downhill acceptance. We only accept improvements, making our algorithm a strict hill climber. Classical SA would occasionally accept a worse solution, with probability decreasing as the temperature drops
Adaptive cooling. Classical SA uses a fixed cooling schedule (e.g., $T_k = T_0 / \log(1+k)$ for the theoretical guarantee). Our schedule only cools when an improvement is found, which adapts the exploration rate to the difficulty of the landscape

Despite these simplifications, our algorithm captures SA's core idea: start with large moves (exploration) and gradually transition to small moves (exploitation). For low-dimensional problems like our 4-parameter CartPole policy, this simplified variant works as well as the full SA.

Theoretical Guarantees

Kirkpatrick et al. proved that SA with logarithmic cooling ( $T_k = c / \log(1+k)$ ) converges to the global optimum in probability. However, this schedule is impractically slow for real problems. In practice, faster geometric schedules ( $T_{k+1} = \alpha T_k$ ) are used, sacrificing the global optimality guarantee for practical convergence speed.

"There is a deep and useful connection between statistical mechanics [...] and multivariate or combinatorial optimization. [...] We have applied this framework to the design of computer hardware, to a specific and practical problem in computer layout."
Kirkpatrick, Gelatt, and Vecchi (1983)

Try It Yourself

The interactive notebook includes exercises:

Decay rate sweep: Try decay values of 0.8, 0.9, 0.95, and 0.99. How does the cooling speed affect convergence? Is there a sweet spot?
True simulated annealing: Modify the algorithm to accept worse solutions with probability $e^{-\Delta / T}$ where $\Delta$ is the score difference and $T$ decays on a fixed schedule. Does it help on CartPole? When would it matter?
Seed sensitivity: Run the algorithm 20 times with different random seeds. What fraction of runs reach 500? How does this compare to CEM's reliability?
Harder environments: Try SA on Acrobot-v1 or MountainCar-v0. Does the 4-parameter linear policy have enough capacity, or do these environments need a richer policy class?

Interactive Tools

Q-Learning Visualiser — Compare SA's derivative-free approach with value-based RL on grid worlds

The Cross-Entropy Method: Solving RL Without Gradients - The population-based companion to SA. Both are derivative-free, but CEM trades sample efficiency for robustness by maintaining 200 candidates per iteration.
MCMC Island Hopping: Understanding Metropolis-Hastings - The acceptance criterion that powers true SA comes directly from the Metropolis algorithm. In MCMC we sample from a distribution; in SA we find its peak.
Genetic Algorithms: From Line Fitting to the Travelling Salesman - Another derivative-free optimisation family. GAs use crossover and mutation on a population; SA uses perturbation on a single solution.

Frequently Asked Questions

How is simulated annealing different from random search?

Random search samples a completely new policy each iteration and tracks the best one found, but it can never refine a promising solution. Simulated annealing builds on previous improvements by perturbing the current best parameters with decreasing noise. This ability to make small refinements to an already-good solution is what pushes SA from "decent" to "perfect" on CartPole.

Why does the algorithm evaluate each candidate over 10 episodes instead of 1?

CartPole has stochastic initial conditions, so a single episode can be misleading. A policy might score 500 on one lucky initialisation and 50 on the next. Averaging over 10 episodes gives a stable estimate of true quality, preventing the algorithm from accepting a lucky fluke or rejecting a good policy due to bad luck.

Is this true simulated annealing?

Not quite. True simulated annealing occasionally accepts worse solutions with a probability that decreases over time, allowing it to escape local optima. Our implementation is a strict hill climber that only accepts improvements. The "annealing" part refers only to the shrinking step size. For CartPole's smooth 4-parameter landscape, this distinction does not matter, but for problems with many local optima, true SA's downhill acceptance becomes essential.

Why does the step size only shrink when an improvement is found?

This creates an adaptive cooling schedule. If the algorithm keeps finding improvements, the step size decays quickly, focusing the search around the current best. If it gets stuck in a plateau, the step size stays large, maintaining broad exploration. This naturally spends more time exploring when stuck and more time refining when making progress.

When would simulated annealing fail compared to population-based methods?

SA struggles in high-dimensional parameter spaces where a single random perturbation is unlikely to improve all parameters at once. It also fails on multi-modal reward landscapes because, as a strict hill climber, it can only find the nearest peak. Population-based methods like the Cross-Entropy Method or genetic algorithms handle both cases better by maintaining diversity across many candidates simultaneously.

The Cross-Entropy Method: Solving RL Without Gradients

Berkan Sesen — Tue, 21 Apr 2026 08:27:46 +0000

Reinforcement learning has accumulated layers of complexity over the years: value functions, policy gradients, replay buffers, target networks. The Cross-Entropy Method predates all of it. Rubinstein introduced it in 1997 for rare-event simulation, and it turned out to solve simple control tasks with almost no machinery. The entire implementation fits in 50 lines. No gradients, no training loops. Just: sample some parameters, test them, keep the best ones, repeat.

The Cross-Entropy Method (CEM) is the algorithm you reach for when you want results without complexity. It treats the policy's parameters as a black box, maintains a probability distribution over them, and iteratively narrows that distribution toward high-performing regions. No gradients required. By the end of this post, you'll implement CEM from scratch, solve CartPole-v1 with a perfect score, and understand why this "naive" approach works so well on problems with manageable parameter spaces.

Let's Build It

Click the badge to open the interactive notebook:

Here's the complete implementation. We use a linear policy with just 4 parameters (one per observation dimension), and CEM finds the perfect weights:

import numpy as np
import gymnasium as gym

def evaluate_policy(env_name, theta):
    """Run one episode with a linear policy: action = 1 if theta @ obs > 0 else 0."""
    env = gym.make(env_name)
    obs, _ = env.reset()
    total_reward = 0
    done = False
    while not done:
        action = 1 if np.dot(theta, obs) > 0 else 0
        obs, reward, terminated, truncated, _ = env.step(action)
        total_reward += reward
        done = terminated or truncated
    env.close()
    return total_reward

def cem(env_name, n_params, batch_size=200, n_iter=50, elite_frac=0.2,
        initial_std=1.0, extra_std=0.5, std_decay_time=25):
    """Cross-Entropy Method for policy search."""
    n_elite = int(np.round(batch_size * elite_frac))
    th_mean = np.zeros(n_params)
    th_std = np.ones(n_params) * initial_std

    for iteration in range(n_iter):
        # Decaying extra noise (Szita & Lörincz 2006)
        noise_multiplier = max(1.0 - iteration / float(std_decay_time), 0)
        sample_std = np.sqrt(th_std + np.square(extra_std) * noise_multiplier)

        # Sample and evaluate
        thetas = th_mean + sample_std * np.random.randn(batch_size, n_params)
        rewards = np.array([evaluate_policy(env_name, th) for th in thetas])

        # Select elite and refit distribution
        elite_inds = rewards.argsort()[-n_elite:]
        elite_thetas = thetas[elite_inds]
        th_mean = elite_thetas.mean(axis=0)
        th_std = elite_thetas.var(axis=0)

        print(f"Iter {iteration+1:3d} | Mean: {rewards.mean():6.1f} | Max: {rewards.max():.0f}")

    return th_mean

best_theta = cem('CartPole-v1', n_params=4)
# Iter   1 | Mean:   66.8 | Max: 500
# Iter  10 | Mean:  384.0 | Max: 500
# Iter  30 | Mean:  495.2 | Max: 500
# Iter  50 | Mean:  499.1 | Max: 500

The population mean reward climbs from 67 to 499 in 50 iterations. Every single sample in the final batch scores near-perfect. Let's verify with 100 evaluation episodes:

scores = [evaluate_policy('CartPole-v1', best_theta) for _ in range(100)]
print(f"Mean: {np.mean(scores):.0f} ± {np.std(scores):.0f}")
# Mean: 500 ± 0

Perfect score. Four parameters, zero gradients, 50 iterations.

What Just Happened?

CEM works by maintaining a Gaussian distribution over policy parameters and repeatedly narrowing it toward the best-performing region. Each iteration has three steps:

Step 1: Sample

We draw batch_size=200 parameter vectors from a Gaussian:

thetas = th_mean + sample_std * np.random.randn(batch_size, n_params)

Each theta is a candidate policy. In iteration 1, the mean is zeros and the standard deviation is 1.0, so we're sampling random policies.

Step 2: Evaluate and Select

We run each candidate policy on CartPole and rank them by total reward. Then we keep only the top 20% (the "elite" set):

rewards = np.array([evaluate_policy(env_name, th) for th in thetas])
elite_inds = rewards.argsort()[-n_elite:]  # Top 40 out of 200
elite_thetas = thetas[elite_inds]

Step 3: Refit the Distribution

We refit the Gaussian to match the elite samples:

th_mean = elite_thetas.mean(axis=0)
th_std = elite_thetas.var(axis=0)

The new mean moves toward parameters that performed well. The new variance shrinks because the elite samples cluster together. Next iteration, we sample from this tighter distribution, generating better candidates on average.

Watching It Converge

The training curve shows how the population improves:

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(iters, mean_rewards, 'b-', alpha=0.7, label='Population mean')
ax.plot(iters, elite_mean_rewards, 'r-', linewidth=2, label='Elite mean')
ax.plot(iters, max_rewards, 'g--', alpha=0.5, label='Best in batch')
ax.axhline(y=500, color='k', linestyle=':', alpha=0.3, label='Max possible (500)')
ax.set_xlabel('Iteration')
ax.set_ylabel('Total Reward')
ax.legend()

The elite mean hits 500 almost immediately (iteration 2). But the population mean takes longer to catch up because the distribution is still wide. By iteration 30, even randomly sampled policies from the learned distribution score near-perfect.

The Distribution Narrows Over Time

To see this visually, here's how the reward distribution across the 200 samples evolves:

fig, axes = plt.subplots(1, 3, figsize=(12, 4))
for ax, iteration_rewards, title in zip(axes, selected_iterations, titles):
    ax.hist(iteration_rewards, bins=20, color='steelblue', edgecolor='white')
    ax.axvline(x=iteration_rewards.mean(), color='red', linestyle='--')

In iteration 1, most policies fail quickly (reward < 100) with a few lucky ones reaching 500. By iteration 10, the distribution is bimodal: many policies near 500 but some still struggling. By iteration 50, the entire population clusters at 500. The distribution has collapsed onto the solution.

Going Deeper

The Noisy Cross-Entropy Method

The original CEM (Rubinstein 1999) has a failure mode: the variance can collapse to zero too quickly, trapping the search in a local optimum. Szita and Lörincz (2006) fixed this with the "noisy" variant that adds decaying extra variance:

where $Z_t = \max(1 - t / T_{\text{decay}},\; 0)$ decays linearly to zero. Early iterations get extra exploration; later iterations trust the elite variance.

This is exactly what our code does:

noise_multiplier = max(1.0 - iteration / float(std_decay_time), 0)
sample_std = np.sqrt(th_std + np.square(extra_std) * noise_multiplier)

The extra_std=0.5 decays over std_decay_time=25 iterations. After iteration 25, the sampling distribution uses only the elite variance.

Hyperparameters

Parameter	Value	Effect
`batch_size`	200	More samples = better coverage but slower per iteration
`elite_frac`	0.2	Lower = more selective, faster convergence, risk of premature collapse
`initial_std`	1.0	Too low = miss good regions; too high = waste samples on extreme policies
`extra_std`	0.5	Noise injection; 0 = original CEM, >0 = noisy CEM
`std_decay_time`	25	How many iterations before extra noise disappears

The most sensitive parameter is elite_frac. At 0.2 (keep top 40 of 200), we balance exploitation and exploration. Setting it to 0.01 (keep top 2) would converge faster in easy environments but collapse in hard ones.

CEM vs Random Search

Both CEM and random search sample 200 policies per iteration. The difference: random search starts fresh every time, while CEM builds on what worked:

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(iters, cem_mean_rewards, 'b-', linewidth=2, label='CEM (population mean)')
ax.plot(iters, random_mean_rewards, 'r--', linewidth=2, label='Random search (mean)')

Random search averages about 60 reward per iteration, forever. CEM reaches 500 because each iteration's distribution is informed by the last. The "select and refit" loop creates a directed search through parameter space.

CEM vs Policy Gradients vs DQN

How does CEM compare to the gradient-based methods we've covered?

Method	What it optimises	Needs gradients?	Scales to large nets?
CEM	Policy parameters directly	No	Poorly (>1000 params)
REINFORCE	Policy parameters via log-prob gradient	Yes	Yes
DQN	Value function (Q-values)	Yes	Yes
Q-Learning	Value function (Q-table)	No	No (tabular only)

CEM's sweet spot: problems with fewer than ~1000 parameters where you want a simple, parallelisable algorithm. For a 4-parameter linear policy on CartPole, CEM is hard to beat. For a million-parameter Atari network, you need policy gradients or DQN.

When NOT to Use CEM

High-dimensional parameter spaces. CEM samples grow exponentially less effective as dimensions increase. A 1000-parameter network needs enormous batch sizes
Environments with sparse rewards. If most policies score zero (e.g., Montezuma's Revenge), the elite set is just noise
When you need sample efficiency. CEM uses 200 episodes per iteration vs REINFORCE using ~5 episodes per batch. If environment evaluations are expensive, gradient methods win
Continuous action spaces with complex dynamics. CEM with a linear policy can only learn linear decision boundaries. Problems requiring nonlinear policies need either a neural network (large parameter space) or a different algorithm

Connection to Genetic Algorithms

If you read the genetic algorithms post, CEM will feel familiar. Both are population-based, derivative-free optimisation methods. The difference is in how they generate the next population:

Genetic algorithms use crossover and mutation operators on individual solutions
CEM fits a probability distribution to the elite set and samples from it

CEM is sometimes called an "estimation of distribution algorithm" (EDA). Instead of recombining individual solutions, it models the structure of good solutions as a distribution and samples new candidates from that model. For real-valued parameter optimisation, this Gaussian model is often more effective than genetic crossover.

Where This Comes From

The Cross-Entropy Method was introduced by Reuven Rubinstein in his 1999 paper "The Cross-Entropy Method for Combinatorial and Continuous Optimization". The name comes from the original application: minimising the cross-entropy (KL divergence) between a reference distribution and the optimal importance sampling distribution for rare-event simulation.

The Core Idea

Rubinstein's insight was that rare-event estimation and optimisation are essentially the same problem. To estimate $P(S(X) \geq \gamma)$ for a rare threshold $\gamma$ , you need to find a sampling distribution that concentrates on high- $S(X)$ regions. The CE method does this by iteratively:

Drawing samples from the current distribution $f(\cdot;\, v_t)$
Selecting the elite samples (those with $S(X) \geq \gamma_t$ )
Updating the distribution parameters to minimise the KL divergence to the empirical elite distribution

For a Gaussian family, step 3 has a closed-form solution: the mean and variance of the elite samples. This is exactly what our implementation does.

The Formal Algorithm

From Rubinstein and Kroese (2004), the CEM update for a parametric family $\{f(\cdot;\, v)\}$ is:

where $I\{\cdot\}$ is the indicator function selecting elite samples. For a multivariate Gaussian with diagonal covariance, this yields:

The sample mean and variance of the elite set. Elegantly simple.

From Rare Events to Tetris

The method found its way into reinforcement learning through Szita and Lörincz (2006), "Learning Tetris Using the Noisy Cross-Entropy Method". They made two key modifications for RL:

Noisy updates: Adding decaying extra variance to prevent premature convergence (the extra_std parameter in our code)
Direct policy search: Treating the policy's weight vector as the parameter to optimise, with episode return as the objective function

"The noisy cross-entropy method adds a time-decreasing noise term to avoid premature convergence of the variance to zero."
Szita and Lörincz (2006)

Their noisy CEM achieved record-breaking performance on Tetris at the time, outperforming methods that required orders of magnitude more computation. Our implementation follows their variant faithfully, including the linear noise decay schedule described in Section 3 of their paper.

Theoretical Properties

Unlike policy gradient methods, CEM has no convergence guarantees to a local optimum. It is a heuristic. However, it has practical advantages:

Embarrassingly parallel: All 200 evaluations per iteration are independent
No reward shaping needed: Works with any scalar objective, even non-differentiable ones
Robust to noisy evaluations: The elite selection acts as a natural filter

The method's simplicity is also its limitation. As Rubinstein and Kroese note, the Gaussian parametric family assumes the optimal parameter region is unimodal. Multi-modal reward landscapes can trap CEM in a single mode.

Try It Yourself

The interactive notebook includes exercises:

Elite fraction sweep: Try elite_frac values of 0.01, 0.1, 0.2, and 0.5. How does selectivity affect convergence speed and stability?
Noisy vs vanilla CEM: Set extra_std=0 and compare convergence. Does the noisy variant help on CartPole, or only on harder problems?
Neural network policy: Replace the linear policy with a small neural net (8 hidden units). How many CEM iterations does it take to solve CartPole now? At what network size does CEM become impractical?
Different environments: Try CEM on Acrobot-v1 or MountainCar-v0. Which environments does CEM handle well, and which expose its limitations?

Interactive Tools

Q-Learning Visualiser — See value-based RL in action and compare it with the policy search approach of CEM

Genetic Algorithms: From Line Fitting to the Travelling Salesman - Another population-based, derivative-free optimisation method. CEM replaces crossover and mutation with distribution fitting.
Policy Gradients: REINFORCE from Scratch with NumPy - The gradient-based alternative for policy search. Uses backpropagation through the policy, which scales to large networks but requires differentiable objectives.
Deep Q-Networks: Experience Replay and Target Networks - Value-based RL with neural networks. A fundamentally different approach that learns what states are valuable rather than directly searching for good policies.

Frequently Asked Questions

Why is the Cross-Entropy Method called "cross-entropy" if it does not use a loss function?

The name comes from the original application in rare-event simulation, where the algorithm minimises the cross-entropy (KL divergence) between the current sampling distribution and the optimal importance sampling distribution. In the reinforcement learning context, the name persists even though the update reduces to simply computing the mean and variance of the elite samples.

How does CEM compare to random search?

Both methods sample candidate policies each iteration, but random search draws from a fixed distribution every time, while CEM updates its distribution based on the best-performing candidates. This directed search means CEM builds on previous successes, converging to good solutions far faster than random search on problems with structure to exploit.

Can CEM solve problems with continuous action spaces?

CEM can optimise over continuous policy parameters, but the policy itself determines how actions are generated. A linear policy with CEM-optimised weights can only produce binary or discrete decisions. For truly continuous action spaces with complex dynamics, you would need a more expressive policy architecture, which increases the parameter count and makes CEM less practical.

What is the role of the elite fraction hyperparameter?

The elite fraction controls how selective the algorithm is when choosing which candidates inform the next distribution. A smaller fraction (e.g. 0.01) converges faster but risks collapsing onto a local optimum. A larger fraction (e.g. 0.5) explores more broadly but converges more slowly. A value around 0.2 is a common default that balances exploitation and exploration.

Why does the noisy CEM variant add extra variance that decays over time?

Without extra variance, the sampling distribution can collapse to near-zero variance too quickly, trapping the search around a potentially suboptimal solution. The decaying noise keeps exploration alive in early iterations when the algorithm is still uncertain about the best region, then gradually disappears to allow precise convergence in later iterations.

PCR vs PLS: When Fewer Features Beat More

Berkan Sesen — Sun, 19 Apr 2026 15:38:56 +0000

How much should a baseball team pay its players? The 1986 Major League season gives us 263 hitters with 19 statistics each: at-bats, hits, home runs, years played, and more. Predicting salary from performance sounds like a textbook regression problem, but 19 correlated features make it anything but. Throw them all into a linear regression and the model fits the training data beautifully but falls apart on held-out players. The coefficient estimates are wildly unstable, and salary predictions swing by thousands on minor input changes.

The fix is not a fancier model. It is fewer features, chosen more carefully. This post covers two classic strategies for doing exactly that: Principal Component Regression (PCR) and Partial Least Squares (PLS).

By the end, you'll understand how both methods compress correlated features into a handful of components, why PLS typically needs fewer components than PCR, and when each approach is the right tool.

Quick Win: Predict Salaries with 6 Features Instead of 19

Click the badge to open the interactive notebook:

We'll use the classic ISLR Hitters dataset: 263 baseball players with 19 features (at-bats, hits, home runs, years played, etc.) predicting salary in thousands of dollars.

import numpy as np
import pandas as pd
from sklearn.preprocessing import scale
from sklearn.decomposition import PCA
from sklearn.linear_model import LinearRegression
from sklearn.cross_decomposition import PLSRegression
from sklearn.model_selection import KFold, cross_val_score, train_test_split
from sklearn.metrics import mean_squared_error

# Load and prepare the Hitters dataset
df = pd.read_csv(
    'https://raw.githubusercontent.com/selva86/datasets/master/Hitters.csv'
).dropna()
dummies = pd.get_dummies(df[['League', 'Division', 'NewLeague']])
y = df['Salary']
X = pd.concat([
    df.drop(['Salary', 'League', 'Division', 'NewLeague'], axis=1).astype('float64'),
    dummies[['League_N', 'Division_W', 'NewLeague_N']]
], axis=1)

# Train/test split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5, random_state=1)

# PCR: PCA on scaled training data, then regression
pca = PCA()
X_train_pc = pca.fit_transform(scale(X_train))
X_test_pc = pca.transform(scale(X_test))

# Use 10-fold CV to find the best number of components
kf = KFold(n_splits=10, shuffle=True, random_state=1)
regr = LinearRegression()
mse_by_ncomp = []
for k in range(1, X.shape[1] + 1):
    score = -cross_val_score(
        regr, X_train_pc[:, :k], y_train.to_numpy(),
        cv=kf, scoring='neg_mean_squared_error'
    ).mean()
    mse_by_ncomp.append(score)

best_k = np.argmin(mse_by_ncomp) + 1
print(f"Best PCR components: {best_k}")
print(f"CV MSE: {mse_by_ncomp[best_k - 1]:,.0f}")

# Evaluate on test set
regr.fit(X_train_pc[:, :best_k], y_train)
pcr_test_mse = mean_squared_error(y_test, regr.predict(X_test_pc[:, :best_k]))
print(f"PCR test MSE ({best_k} components): {pcr_test_mse:,.0f}")

# Compare to full OLS
regr_full = LinearRegression()
regr_full.fit(X_train_pc, y_train)
ols_test_mse = mean_squared_error(y_test, regr_full.predict(X_test_pc))
print(f"Full OLS test MSE (19 features): {ols_test_mse:,.0f}")

The result: PCR with just 6 components achieves a test MSE of ~112,000, beating full OLS (test MSE ~117,000) using all 19 features. Fewer features, better predictions.

PCR vs PLS: The Key Difference

Now let's try PLS, which uses the target variable during dimension reduction:

# PLS: find the best number of components via CV
pls_mse = []
for k in range(1, X.shape[1] + 1):
    pls = PLSRegression(n_components=k)
    score = -cross_val_score(
        pls, scale(X_train), y_train.to_numpy(),
        cv=kf, scoring='neg_mean_squared_error'
    ).mean()
    pls_mse.append(score)

best_pls_k = np.argmin(pls_mse) + 1
pls_best = PLSRegression(n_components=2)
pls_best.fit(scale(X_train), y_train)
pls_test_mse = mean_squared_error(y_test, pls_best.predict(scale(X_test)))
print(f"PLS test MSE (2 components): {pls_test_mse:,.0f}")

PLS with just 2 components achieves a test MSE of ~105,000, beating both PCR and OLS. That is the power of supervised dimension reduction: PLS finds the directions that matter for the target, not just the directions of maximum variance.

What Just Happened?

Both methods solve the same problem: your 19 features are correlated (career stats like CAtBat, CHits, CRuns all move together), so fitting a separate coefficient for each one leads to noisy, unstable estimates. The solution is to compress correlated features into a smaller set of components before regressing.

The difference is how they choose those components.

PCR: Unsupervised, Then Regress

PCR works in two steps:

PCA finds the directions of maximum variance in $X$ , ignoring $y$ entirely
Linear regression fits $y$ on the top $k$ principal components

# Step 1: PCA finds directions of maximum variance
pca = PCA()
X_train_pc = pca.fit_transform(scale(X_train))  # 19 features → 19 PCs

# Step 2: Regress salary on just the first k PCs
k = 6
regr = LinearRegression()
regr.fit(X_train_pc[:, :k], y_train)

The first principal component captures the direction along which the features vary the most. In our Hitters data, PC1 captures 39.9% of the total variance, and by PC7 we're at 93.4%.

But here's the catch: the directions of maximum variance in $X$ are not necessarily the directions most useful for predicting $y$ . PC1 might capture the spread between high-career and low-career players, but if salary depends more on a subtle interaction between recent performance and league, that signal could be buried in PC8 or PC12.

PLS: Supervised from the Start

PLS finds directions that simultaneously:

Explain variance in $X$ (like PCA)
Correlate with $y$ (unlike PCA)

# PLS finds directions that maximise covariance between X and y
pls = PLSRegression(n_components=2)
pls.fit(scale(X_train), y_train)
predictions = pls.predict(scale(X_test))

This is why PLS needs only 2 components where PCR needs 6. PLS searches directly for the features that predict salary, while PCR has to hope that the high-variance directions in $X$ also happen to predict $y$ .

Choosing the Number of Components

Both methods use 10-fold cross-validation to select the number of components:

import matplotlib.pyplot as plt

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

# PCR
ax1.plot(range(1, 20), mse_by_ncomp, '-o', color='steelblue', markersize=5)
ax1.axvline(x=best_k, color='red', linestyle='--', alpha=0.7, label=f'Best: {best_k}')
ax1.set_xlabel('Number of Components')
ax1.set_ylabel('10-Fold CV MSE')
ax1.set_title('PCR')
ax1.legend()
ax1.grid(True, alpha=0.3)

# PLS
ax2.plot(range(1, 20), pls_mse, '-s', color='darkorange', markersize=5)
ax2.axvline(x=2, color='green', linestyle='--', alpha=0.7, label='Selected: 2 (parsimonious)')
ax2.axvline(x=best_pls_k, color='red', linestyle=':', alpha=0.5, label=f'CV min: {best_pls_k}')
ax2.set_xlabel('Number of Components')
ax2.set_ylabel('10-Fold CV MSE')
ax2.set_title('PLS')
ax2.legend(fontsize=9)
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

The PCR curve dips at 6 components and rises again: adding noisy components hurts predictions. The PLS curve is more interesting: the strict CV minimum is at 11 components, but 2 components achieve nearly the same MSE (143,564 vs 142,554). We select 2 because the simpler model generalises better on the test set (MSE 104,839 vs 106,891 with 11). This is a common pattern: when the CV curve is flat near the minimum, prefer the simpler model.

A Peek Inside the Components

What do these principal components actually capture? The PCA loadings reveal which original features contribute to each component:

loadings = pd.DataFrame(
    pca.components_[:5].T,
    columns=[f'PC{i+1}' for i in range(5)],
    index=X.columns
)
print(loadings.sort_values('PC1', ascending=False).round(3))

The heatmap reveals the correlation structure clearly. PC1 (39.9% of variance) is dominated by career statistics: CRuns, CRBI, CHits, CAtBat, and CHmRun all have loadings above 0.30. PC2 (21.5%) separates current-season stats (AtBat, Hits, Runs with positive loadings) from career longevity (Years with a negative loading). PC3 picks up the league indicator variables. PCA compresses these correlated groups into single components, which is exactly why dimension reduction works here.

Going Deeper

The Mathematics

PCR decomposes $X$ using PCA:

where $V$ contains the principal component directions (eigenvectors of $X^TX$ ). We keep only the first $k$ columns of $Z = XV_k$ and regress:

This is equivalent to OLS on the reduced feature set. The key insight: since the PCs are orthogonal, the regression coefficients don't change when you add or remove components. Each component's contribution is independent.

PLS maximises the covariance between $X$ and $y$ :

The first PLS direction $w_1$ is simply $X^T y$ (normalised): the covariance between each feature and the target. Subsequent directions are found by deflating $X$ and repeating.

When PCR Wins, When PLS Wins

PCR is better when:

The high-variance directions in $X$ genuinely predict $y$ (common in spectroscopy, genomics)
You have many features and few observations (PCA provides stable variance estimates)
You want an unsupervised feature extraction that you can reuse across multiple targets

PLS is better when:

The predictive signal sits in low-variance directions of $X$
You have a single target and want the most efficient compression
Your features include many irrelevant high-variance variables

In our Hitters example, PLS wins convincingly: 2 components vs 6, and lower test error. The salary signal does not align perfectly with the directions of maximum variance in the batting statistics.

The Bias-Variance Tradeoff

Both methods trade bias for lower variance:

Method	Components	Test MSE	RMSE ($k)
Full OLS	19	117,301	342
PCR	6	112,167	335
PLS	2	104,839	324
Ridge	--	99,741	316

Full OLS uses all 19 features but has high variance (unstable coefficients). PCR and PLS introduce some bias by discarding information, but the reduction in variance more than compensates. Ridge regression (included for comparison) achieves the lowest error by shrinking coefficients rather than discarding components.

The practical message: when features are correlated, you rarely need all of them. The question is whether to reduce dimensions unsupervised (PCR), supervised (PLS), or regularise without reducing dimensions at all (Ridge).

When NOT to Use PCR or PLS

Few features, many observations. If $p \ll n$ , multicollinearity is less of a problem and OLS works fine.
Interpretability is critical. The principal components are linear combinations of all features, so individual feature effects are obscured. If you need to say "an extra home run is worth $X in salary," use Ridge or Lasso instead.
Non-linear relationships. PCR and PLS are linear methods. For non-linear patterns, consider Gaussian process regression or tree-based models.
Sparse signals. If only a few features matter and the rest are noise, Lasso (L1 regularisation) does feature selection rather than feature combination, which is usually more effective.

Deep Dive: The Papers

Principal Component Regression

The idea of using principal components as regression predictors dates to Massy (1965), "Principal Components Regression in Exploratory Statistical Research", published in the Journal of the American Statistical Association.

Massy was working on marketing research problems where survey data had dozens of correlated variables. He proposed a two-step procedure: extract principal components, then regress on the top $k$ . His key insight:

"By using the principal components as the independent variables in the regression, we avoid the multicollinearity problem since the components are orthogonal."

The underlying PCA dates back further to Hotelling (1933), "Analysis of a complex of statistical variables into principal components", Journal of Educational Psychology. Hotelling formalised the eigenvalue decomposition of the covariance matrix, though the core idea appeared even earlier in Pearson (1901).

Partial Least Squares

PLS was developed by Herman Wold in the 1960s and 1970s, originally for econometrics. The foundational paper is Wold (1975), "Soft modelling by latent variables: the non-linear iterative partial least squares (NIPALS) approach," in Perspectives in Probability and Statistics.

Herman's son, Svante Wold, later popularised PLS in chemometrics with a landmark review: Wold, Sjostrom & Eriksson (2001), "PLS-regression: a basic tool of chemometrics", Chemometrics and Intelligent Laboratory Systems.

The modern computational algorithm used in most implementations (including sklearn) is SIMPLS by de Jong (1993), "SIMPLS: An alternative approach to partial least squares regression". de Jong's algorithm computes PLS components without the iterative deflation step, making it both faster and numerically more stable.

The ISLR Connection

This tutorial is based on the lab exercise in James, Witten, Hastie & Tibshirani (2021), An Introduction to Statistical Learning, Chapter 6. ISLR provides an excellent treatment of PCR and PLS in the context of the bias-variance tradeoff, alongside Ridge and Lasso regression.

The Hitters dataset used here has become a standard benchmark for comparing regularisation and dimension reduction methods. With 19 correlated features, it sits in the sweet spot where these methods make a visible difference.

Try It Yourself

The interactive notebook includes exercises:

Scree plot. Plot the explained variance per component and the cumulative curve. How many components do you need to capture 95% of the variance?
PLS loadings. Compare the PLS weight vectors to the PCA loadings. Which features does PLS prioritise that PCA does not?
Ridge vs PCR. Add a Ridge regression (with RidgeCV) to the comparison. In what sense is Ridge a "soft" version of PCR?
Log-transform the target. Salary is right-skewed. Does predicting $\log(\text{Salary})$ change which method wins?

Understanding PCR and PLS builds directly on linear regression and connects to Bayesian inference through the regularisation-as-prior interpretation. When the linear assumption breaks down, Gaussian process regression offers a non-parametric alternative that handles high-dimensional inputs gracefully.

Interactive Tools

Regression Playground — Fit and compare regression models interactively in the browser

Linear Regression Five Ways — The foundation both PCR and PLS build on
LDA vs PCA: Supervised vs Unsupervised Dimensionality Reduction — The classification counterpart to PCR vs PLS
From MLE to Bayesian Inference — How regularisation connects to Bayesian priors
Gaussian Process Regression from Scratch — A non-parametric alternative when linearity breaks down

Frequently Asked Questions

What is the key difference between PCR and PLS?

PCR finds directions of maximum variance in the features without considering the target variable, then regresses on those directions. PLS finds directions that maximise the covariance between the features and the target simultaneously. Because PLS is supervised from the start, it typically needs fewer components to achieve the same predictive performance.

When should I use PCR instead of PLS?

PCR is preferable when the high-variance directions in your features genuinely predict the target, which is common in spectroscopy and genomics. It is also useful when you want an unsupervised feature extraction that can be reused across multiple target variables. PLS is better when the predictive signal sits in low-variance directions or when many high-variance features are irrelevant to the outcome.

How do I choose the number of components for PCR or PLS?

Use k-fold cross-validation to evaluate predictive performance at each number of components and select the value that minimises the cross-validation error. When the error curve is flat near the minimum, prefer the simpler model with fewer components, as it tends to generalise better on unseen data.

Why did PLS with 2 components beat PCR with 6 components on the Hitters dataset?

The salary signal in the baseball data does not align well with the directions of maximum variance in the batting statistics. Career statistics dominate the first few principal components, but salary depends on a subtler combination of recent performance and league factors. PLS finds these salary-relevant directions directly, so it needs far fewer components.

How does PCR compare to Ridge regression?

Both methods address multicollinearity, but in different ways. PCR discards the least important principal components entirely, introducing a hard cutoff. Ridge regression shrinks all coefficients towards zero without discarding any, acting as a soft version of dimension reduction. Ridge often achieves lower test error because it retains some information from every direction.

Can I interpret individual feature effects with PCR or PLS?

Not directly. The components are linear combinations of all original features, so individual feature effects are obscured. If you need to say that a specific feature is worth a certain amount, use Ridge or Lasso regression instead, which produce interpretable coefficients for each original variable.

Text Classification from Scratch: TF-IDF and Naive Bayes

Berkan Sesen — Fri, 17 Apr 2026 12:46:47 +0000

Every morning, your inbox separates spam from real email. News apps sort articles into sports, tech, and politics. Customer support systems route tickets to the right team. Behind all of these is text classification: teaching a machine to read a document and assign it a category.

The building blocks are simpler than you might expect. You need a way to convert text into numbers (TF-IDF), a classifier that works well with sparse, high-dimensional data (Naive Bayes), and a few lines of code to tie them together. No deep learning, no GPUs, no embeddings.

By the end of this post, you'll classify news articles into 20 categories with 77% accuracy using just 10 lines of Python, then push that to 84% with hyperparameter tuning. You'll understand exactly how TF-IDF works and why the "naive" independence assumption in Naive Bayes is a feature, not a bug.

Let's Build It

Click the badge to open the interactive notebook:

Here's the complete classifier. We use scikit-learn's 20 Newsgroups dataset, which contains around 18,000 posts across 20 topics, from computer graphics to religion to space exploration:

from sklearn.datasets import fetch_20newsgroups
from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import Pipeline
from sklearn.metrics import accuracy_score

# Load training and test data
twenty_train = fetch_20newsgroups(subset='train', shuffle=True, random_state=42)
twenty_test = fetch_20newsgroups(subset='test', shuffle=True, random_state=42)

# Build the pipeline: raw text → word counts → TF-IDF → Naive Bayes
text_clf = Pipeline([
    ('vect', CountVectorizer()),
    ('tfidf', TfidfTransformer()),
    ('clf', MultinomialNB()),
])

# Train and evaluate
text_clf.fit(twenty_train.data, twenty_train.target)
predicted = text_clf.predict(twenty_test.data)
print(f'Accuracy: {accuracy_score(twenty_test.target, predicted):.1%}')
# Accuracy: 77.4%

With 10 lines of modelling code, we classify documents into one of 20 categories at 77.4% accuracy on unseen data. Random guessing would give 5%.

Let's test it on fresh sentences the model has never seen:

docs_new = [
    'OpenGL shading techniques for real-time rendering',
    'The Detroit Tigers signed a new pitcher today',
    'NASA launched the James Webb telescope last year',
    'Is there evidence for the existence of God?',
]

predicted_new = text_clf.predict(docs_new)
for doc, category in zip(docs_new, predicted_new):
    print(f'{twenty_train.target_names[category]:>28s}  ←  {doc}')

            comp.graphics  ←  OpenGL shading techniques for real-time rendering
        rec.sport.baseball  ←  The Detroit Tigers signed a new pitcher today
                 sci.space  ←  NASA launched the James Webb telescope last year
    soc.religion.christian  ←  Is there evidence for the existence of God?

The model correctly identifies the topic of each sentence. It works by finding which words are most characteristic of each category.

The confusion matrix reveals where the classifier struggles. Related categories like comp.sys.ibm.pc.hardware and comp.sys.mac.hardware (both about computer hardware) are frequently confused, as are talk.religion.misc and soc.religion.christian. These make intuitive sense: documents about Mac hardware and PC hardware use very similar vocabulary.

What Just Happened?

Three components work in sequence: CountVectorizer turns text into word counts, TfidfTransformer re-weights those counts to highlight distinctive words, and MultinomialNB learns which words signal which categories.

Step 1: Turning Text into Numbers

A machine learning model can't read English. It needs numbers. The simplest conversion is the bag of words: count how many times each word appears in a document, ignoring order entirely.

from sklearn.feature_extraction.text import CountVectorizer

corpus = [
    'The cat sat on the mat',
    'The dog sat on the log',
    'The cat chased the dog',
]
vectorizer = CountVectorizer()
X = vectorizer.fit_transform(corpus)

print(vectorizer.get_feature_names_out())
# ['cat', 'chased', 'dog', 'log', 'mat', 'on', 'sat', 'the']

print(X.toarray())
# [[1, 0, 0, 0, 1, 1, 1, 2],
#  [0, 0, 1, 1, 0, 1, 1, 2],
#  [1, 1, 1, 0, 0, 0, 0, 2]]

Each row is a document. Each column is a word from the vocabulary. The value is the word count. Notice that "the" always gets a count of 2, regardless of the document. It's everywhere, so it carries no information about which document you're looking at.

On the 20 Newsgroups training set, CountVectorizer discovers around 130,000 unique tokens. Each document becomes a vector of 130,000 dimensions, mostly zeros (since any single post uses only a tiny fraction of the full vocabulary).

Step 2: Weighting Words That Matter

Not all words are equally informative. Words like "the", "is", and "a" appear in every document. What we want are words that are common within a specific category but rare overall. This is exactly what TF-IDF (Term Frequency, Inverse Document Frequency) captures.

The weight for word $t$ in document $d$ is:

Where:

TF (term frequency) = how often the word appears in this document
IDF (inverse document frequency) = $\log\!\frac{1+N}{1+n_t}+1$ , where $N$ is the total number of documents and $n_t$ is the number of documents containing word $t$

A word that appears in every document gets a low IDF, shrinking its weight. A word that appears in only a few documents gets a high IDF, amplifying its signal.

import numpy as np
from sklearn.feature_extraction.text import TfidfTransformer

tfidf = TfidfTransformer()
X_tfidf = tfidf.fit_transform(X)
print(np.round(X_tfidf.toarray(), 2))

After TF-IDF weighting, the document vectors highlight what's distinctive about each text rather than what's common across all of them.

Step 3: Naive Bayes Classification

Naive Bayes applies Bayes' theorem to classify documents. Given a document with words $w_1, w_2, \ldots, w_n$ , it computes:

The "naive" part is the assumption that words are conditionally independent given the class. This is obviously wrong: the word "neural" is far more likely to appear near "network" than near "baseball". But the simplification works remarkably well in practice because:

We only need the ranking right, not the exact probabilities. If $P(\text{sci.space} \mid \text{doc})$ is the highest, the prediction is correct even if the probability value itself is off.
Independence errors tend to cancel out across thousands of features.
The alternative (modelling all word dependencies) is intractable for vocabularies of 130,000 words.

The MultinomialNB variant uses word counts (or TF-IDF weights) as features and models $P(w_i \mid \text{class})$ as a multinomial distribution. The parameters are estimated via maximum likelihood: the probability of word $w_i$ in class $c$ is simply the fraction of times $w_i$ appears in training documents of class $c$ , with Laplace smoothing to handle words never seen in training.

The Pipeline: Composing the Steps

Scikit-learn's Pipeline chains these three transformations so you can treat the entire workflow as a single estimator:

text_clf = Pipeline([
    ('vect', CountVectorizer()),     # raw text → word counts
    ('tfidf', TfidfTransformer()),   # word counts → TF-IDF weights
    ('clf', MultinomialNB()),        # TF-IDF vectors → class predictions
])

When you call text_clf.fit(X, y), it runs CountVectorizer.fit_transform(), feeds the output to TfidfTransformer.fit_transform(), then passes the result to MultinomialNB.fit(). At prediction time, the same chain runs in sequence. This also means you can do grid search over any parameter in the pipeline using the double-underscore naming convention (vect__ngram_range, clf__alpha).

Going Deeper

Beating the Baseline

Naive Bayes at 77.4% is a strong starting point, but we can improve it in three ways: removing noise (stop words), capturing phrases (bigrams), and tuning the smoothing parameter.

Stop words are common words ("the", "is", "at") that carry little discriminative value. Removing them reduces noise and bumps accuracy from 77.4% to 81.7%:

text_clf_stop = Pipeline([
    ('vect', CountVectorizer(stop_words='english')),
    ('tfidf', TfidfTransformer()),
    ('clf', MultinomialNB()),
])
text_clf_stop.fit(twenty_train.data, twenty_train.target)
print(f'NB + stop words: {accuracy_score(twenty_test.target, text_clf_stop.predict(twenty_test.data)):.1%}')
# NB + stop words: 81.7%

A 4-point gain for one parameter change.

Grid search systematically explores combinations of pipeline parameters. The naming convention (vect__, tfidf__, clf__) lets you reach into any pipeline step:

from sklearn.model_selection import GridSearchCV

parameters = {
    'vect__ngram_range': [(1, 1), (1, 2)],  # unigrams vs unigrams+bigrams
    'tfidf__use_idf': (True, False),         # use IDF weighting or not
    'clf__alpha': (1e-2, 1e-3),              # smoothing strength
}

gs_clf = GridSearchCV(text_clf, parameters, cv=5, n_jobs=-1)
gs_clf.fit(twenty_train.data, twenty_train.target)

print(f'Best CV score: {gs_clf.best_score_:.1%}')
print(f'Best params: {gs_clf.best_params_}')
print(f'Test accuracy: {accuracy_score(twenty_test.target, gs_clf.predict(twenty_test.data)):.1%}')
# Best CV score: 91.6%
# Best params: {'clf__alpha': 0.001, 'tfidf__use_idf': True, 'vect__ngram_range': (1, 2)}
# Test accuracy: 83.6%

The best configuration uses bigrams (ngram_range=(1,2)), IDF weighting, and weak smoothing (alpha=0.001). Bigrams capture phrases like "White House" or "hard drive" that individual words miss. The 5-fold CV score (91.6%) is higher than the test accuracy (83.6%) because cross-validation evaluates on data drawn from the same distribution as training, while the test set may contain authors, topics, or writing styles not seen during training.

If you've read our hyperparameter optimisation post, you'll recognise grid search as the brute-force baseline. With only 8 combinations to evaluate here, it's fast enough.

SVM: A Stronger Classifier

Swapping Naive Bayes for a linear SVM (support vector machine) gives a larger improvement than any amount of NB tuning:

from sklearn.linear_model import SGDClassifier

text_clf_svm = Pipeline([
    ('vect', CountVectorizer()),
    ('tfidf', TfidfTransformer()),
    ('clf-svm', SGDClassifier(loss='hinge', penalty='l2',
                               alpha=1e-3, max_iter=100,
                               random_state=42)),
])
text_clf_svm.fit(twenty_train.data, twenty_train.target)
print(f'SVM accuracy: {accuracy_score(twenty_test.target, text_clf_svm.predict(twenty_test.data)):.1%}')
# SVM accuracy: 82.4%

That's 82.4% out of the box, without any tuning. Grid search for SVM yields 83.5%, virtually identical to the tuned Naive Bayes.

The story is clear: the biggest gains come from better feature representation (bigrams, stop word removal, IDF weighting) rather than the choice of classifier. With good features, even the "naive" model performs competitively.

What the Model Actually Learns

What words does the classifier rely on? Raw class-conditional probabilities are dominated by common words like "the" and "of". To find truly discriminative features, we compare each word's log-probability within a class against its average across all classes:

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer

tfidf_vect = TfidfVectorizer(stop_words='english', max_df=0.9, min_df=5)
X_tfidf = tfidf_vect.fit_transform(twenty_train.data)
clf_disc = MultinomialNB().fit(X_tfidf, twenty_train.target)

feature_names = np.array(tfidf_vect.get_feature_names_out())
log_probs = clf_disc.feature_log_prob_
mean_log_prob = np.mean(log_probs, axis=0)
discriminativeness = log_probs - mean_log_prob

for i, category in enumerate(twenty_train.target_names):
    top_indices = discriminativeness[i].argsort()[-5:][::-1]
    print(f'{category}: {", ".join(feature_names[top_indices])}')

The model learns sensible patterns. sci.space relies on words like "space", "orbit", and "nasa". rec.sport.baseball relies on "baseball", "team", and "pitching". talk.politics.mideast picks up "israel", "armenian", and "turkish". These are the words that carry the strongest evidence for each category, well beyond their background frequency.

Stemming: Reducing Words to Roots

Stemming maps words to their root form ("running" to "run", "computers" to "comput"). This merges related word forms into a single feature, reducing vocabulary size:

import nltk
from nltk.stem.snowball import SnowballStemmer

nltk.download('punkt', quiet=True)
stemmer = SnowballStemmer('english', ignore_stopwords=True)

class StemmedCountVectorizer(CountVectorizer):
    def build_analyzer(self):
        analyzer = super().build_analyzer()
        return lambda doc: [stemmer.stem(w) for w in analyzer(doc)]

text_clf_stemmed = Pipeline([
    ('vect', StemmedCountVectorizer(stop_words='english')),
    ('tfidf', TfidfTransformer()),
    ('clf', MultinomialNB(fit_prior=False)),
])
text_clf_stemmed.fit(twenty_train.data, twenty_train.target)
print(f'NB + stemming + stop words: '
      f'{accuracy_score(twenty_test.target, text_clf_stemmed.predict(twenty_test.data)):.1%}')

Stemming often gives a small additional boost. The original code uses the Snowball stemmer, a refined version of Porter's classic 1980 algorithm that handles irregular forms more gracefully.

When NOT to Use Bag-of-Words

This approach has clear limitations:

Word order is lost. "Dog bites man" and "man bites dog" produce the same vector. For tasks where order matters (sentiment analysis, textual entailment), you need sequence models or contextual embeddings.
Synonyms are invisible. If test documents use different words for the same concepts, they won't match. Pre-trained embeddings (Word2Vec, BERT) capture semantic similarity.
Short documents suffer. With only a few words, the sparse vector is too noisy for reliable classification. Transformer models handle short texts much better.
Scalability ceiling. As the number of overlapping categories grows, the independence assumption becomes more costly.

For many practical applications, TF-IDF with Naive Bayes remains hard to beat when you factor in the ratio of performance to complexity. It trains in seconds, requires no GPU, and produces interpretable results.

Where This Comes From

McCallum & Nigam (1998)

The foundational paper for Naive Bayes text classification is McCallum, A. & Nigam, K. (1998) "A Comparison of Event Models for Naive Bayes Text Classification", presented at the AAAI Workshop on Learning for Text Categorization.

They compared two Naive Bayes variants for text:

Multi-variate Bernoulli: each word is a binary feature (present or absent). This is BernoulliNB in scikit-learn.
Multinomial: each word is a count feature. This is the MultinomialNB our pipeline uses.

"We find that the multinomial model is almost uniformly superior, especially for large vocabulary sizes."

The multinomial model works better because it uses word frequency information. A document mentioning "baseball" 15 times is stronger evidence for rec.sport.baseball than one mentioning it once. The Bernoulli model discards this frequency signal entirely.

The Multinomial Model

Formally, the predicted class for a document $d$ is:

Where:

$P(c)$ is the class prior (fraction of training documents in class $c$ )
$n_i(d)$ is the count of word $w_i$ in document $d$
$P(w_i \mid c)$ is estimated with Laplace smoothing:

The smoothing parameter $\alpha$ prevents zero probabilities for words that never appeared in a particular class during training. Our grid search found $\alpha = 0.001$ optimal, meaning the model trusts the training data more and smooths less aggressively than the default $\alpha = 1.0$ .

TF-IDF: Salton & Buckley (1988)

TF-IDF was formalised by Salton, G. & Buckley, C. (1988) "Term-weighting approaches in automatic text retrieval", Information Processing & Management. The core idea predates this work: Sparck Jones proposed inverse document frequency in 1972.

Scikit-learn's variant uses:

The "+1" terms prevent division by zero and ensure no word gets zero weight. After computing TF-IDF, each document vector is L2-normalised to unit length.

Historical Context

Text classification has a long lineage:

Maron (1961) — First automatic text classification using probabilistic indexing
Salton (1971) — The SMART retrieval system, introducing many weighting schemes
Sparck Jones (1972) — Inverse document frequency
Lewis (1998) — The Reuters benchmark that standardised evaluation
Joachims (1998) — Showed SVMs outperform NB on text (our results confirm this: 82.4% vs 77.4%)
McCallum & Nigam (1998) — Systematic comparison of NB event models

Today, transformer-based models (BERT, GPT) dominate text classification benchmarks. But TF-IDF with Naive Bayes remains the standard baseline for its speed, interpretability, and surprising competitiveness.

Try It Yourself

The interactive notebook includes exercises:

Subset classification — Use only 4 categories (comp.graphics, rec.sport.baseball, sci.space, talk.politics.mideast). How much does accuracy improve with fewer, more distinct categories?
Feature engineering — Add min_df=5 and max_df=0.5 to CountVectorizer to trim rare and ubiquitous words. How does this affect accuracy and vocabulary size?
Bernoulli vs Multinomial — Replace MultinomialNB with BernoulliNB. Does the McCallum & Nigam finding hold on this dataset?
Beyond bag-of-words — Use TfidfVectorizer with sublinear_tf=True and character n-grams (analyzer='char_wb', ngram_range=(3,5)). Character n-grams capture morphological patterns that word-level features miss.

Interactive Tools

Classification Metrics Calculator — Compute precision, recall, F1, and other metrics from your own confusion matrix
Bayes' Theorem Calculator — Explore the Bayesian reasoning that underpins Naive Bayes classification

Maximum Likelihood Estimation from Scratch — The estimation method behind Naive Bayes parameter learning
From MLE to Bayesian Inference — The Bayes' theorem foundation that powers Naive Bayes classification
Hyperparameter Optimization: Grid vs Random vs Bayesian — A deeper look at grid search and smarter alternatives

Frequently Asked Questions

Why is Naive Bayes called "naive"?

The "naive" refers to the conditional independence assumption: the model assumes that each word in a document is independent of every other word, given the class. This is clearly wrong (e.g. "neural" and "network" tend to co-occur), but it works surprisingly well in practice because classification only requires getting the ranking of class probabilities right, not the exact values. Independence errors tend to cancel out across thousands of features.

What is the difference between TF-IDF and raw word counts?

Raw word counts treat all words equally, so common words like "the" and "is" dominate the representation despite carrying no discriminative information. TF-IDF re-weights each word by how rare it is across the entire corpus. Words that appear in many documents get downweighted, while words distinctive to a few documents get amplified. This makes the representation much more informative for classification.

When should I use Naive Bayes instead of a transformer model like BERT?

Naive Bayes with TF-IDF is an excellent choice when you need fast training (seconds, not hours), interpretability (you can inspect which words drive predictions), or when labelled data is limited. It also requires no GPU. For tasks where word order matters (sentiment analysis, entailment) or where you need state-of-the-art accuracy on competitive benchmarks, transformer models will outperform it significantly.

What does the smoothing parameter alpha do in MultinomialNB?

Alpha controls Laplace smoothing, which prevents zero probabilities for words that never appeared in a particular class during training. With alpha = 1.0 (the default), the model adds a pseudocount of 1 to every word-class combination. Smaller values like 0.001 trust the training data more and smooth less aggressively. The optimal value depends on your dataset and can be found through cross-validation.

Why does the model confuse related categories like PC hardware and Mac hardware?

The bag-of-words representation captures which words appear in a document but not the subtle semantic differences between closely related topics. Categories like PC hardware and Mac hardware share a large portion of their vocabulary (words like "drive", "memory", "board", "system"). The model can only distinguish them by the few words unique to each category, which may not always be present in a given document.

Can TF-IDF handle languages other than English?

Yes. TF-IDF is language-agnostic at its core since it operates on tokens, not linguistic structures. However, you may need to adjust tokenisation for languages without clear word boundaries (e.g. Chinese or Japanese) and consider language-specific stop word lists. Stemming and lemmatisation tools are also language-dependent, so you would need appropriate resources for your target language.

AI Experts Are Dead. Long Live the AI Experts.

Berkan Sesen — Wed, 15 Apr 2026 07:51:05 +0000

Last month, my eight-year-old built a Flappy Bird clone from scratch. He can't really type yet. He certainly can't write Python. What he can do is talk to Claude while I whisper in his ear what to say next. Within an hour, he had a working game: a bird, pipes, a score counter, gravity. He's an "AI expert" now.

And honestly? So is your dentist, your cousin's teenager, and the recruiter who just messaged you on LinkedIn. The barrier to "using AI" has collapsed to the cost of typing a sentence in English. This is genuinely wonderful. Democratisation of powerful technology is how we got the internet, smartphones, and open-source software.

But there's an asymmetry hiding behind this accessibility: while using AI has never been cheaper, building AI has never been more expensive. Training GPT-4 cost over $100 million. Llama 3 required 24,000 GPUs running for months. The companies that can afford to train foundation models from scratch fit comfortably in a single conference room. We've democratised the interface and monopolised the engine.

So where does that leave the engineers, the domain experts, the people who actually know things about medicine, law, finance, or logistics? Somewhere in between. And that somewhere has a name: fine-tuning. For a few hundred to a few thousand dollars, you can take a foundation model and make it yours, trained on your data, speaking your domain's language, following your formatting rules. Not building the engine from scratch, but tuning it to your track.

By the end of this post, you'll fine-tune a model on Azure OpenAI, understand the LoRA algorithm that makes it computationally feasible, and know exactly where fine-tuning sits in the hierarchy from prompt engineering to pre-training.

Quick Win: Fine-Tune on Azure in 20 Lines

Let's start with the punchline. Here's everything you need to fine-tune a GPT model on Azure OpenAI. No Colab badge here (you'll need Azure credentials), but the code itself is almost disappointingly simple.

Prepare Your Training Data

Azure expects JSONL format: one JSON object per line, each containing a conversation. Here's what training data looks like for a medical coding assistant:

{"messages": [{"role": "system", "content": "You are a medical coding assistant. Map clinical descriptions to ICD-10 codes."}, {"role": "user", "content": "Patient presents with acute bronchitis"}, {"role": "assistant", "content": "J20.9 — Acute bronchitis, unspecified"}]}
{"messages": [{"role": "system", "content": "You are a medical coding assistant. Map clinical descriptions to ICD-10 codes."}, {"role": "user", "content": "Type 2 diabetes with diabetic chronic kidney disease"}, {"role": "assistant", "content": "E11.22 — Type 2 diabetes mellitus with diabetic chronic kidney disease"}]}
{"messages": [{"role": "system", "content": "You are a medical coding assistant. Map clinical descriptions to ICD-10 codes."}, {"role": "user", "content": "Essential hypertension"}, {"role": "assistant", "content": "I10 — Essential (primary) hypertension"}]}

Each line is a complete conversation with a system prompt, user input, and the desired assistant response.

Upload Data and Launch Fine-Tuning

from openai import AzureOpenAI

client = AzureOpenAI(
    azure_endpoint="https://YOUR_RESOURCE.openai.azure.com/",
    api_key="YOUR_API_KEY",
    api_version="2025-03-01-preview",
)

# Upload training file
training_file = client.files.create(
    file=open("training_data.jsonl", "rb"),
    purpose="fine-tune",
)

# Create fine-tuning job
job = client.fine_tuning.jobs.create(
    training_file=training_file.id,
    model="gpt-4o-mini-2024-07-18",
)

print(f"Job ID: {job.id}")
print(f"Status: {job.status}")

Check Status and Use Your Model

# Check progress
job = client.fine_tuning.jobs.retrieve(job.id)
print(f"Status: {job.status}")  # queued → running → succeeded

# Once succeeded, deploy and use your model
if job.status == "succeeded":
    response = client.chat.completions.create(
        model=job.fine_tuned_model,
        messages=[
            {"role": "system", "content": "You are a medical coding assistant."},
            {"role": "user", "content": "Chronic obstructive pulmonary disease with acute exacerbation"},
        ],
    )
    print(response.choices[0].message.content)
    # J44.1 — COPD with (acute) exacerbation

That's it. Azure handles the infrastructure, the training loop, the checkpointing, and the deployment. You provide the data; it returns a model that speaks your domain.

What Just Happened?

Behind those 20 lines, a lot happened. Let's unpack it.

The JSONL Format

Each training example is a conversation in the chat completions format you already know. The key fields:

role: "system": Sets the persona. Keep this consistent across examples.
role: "user": The input your model will receive in production.
role: "assistant": The exact output you want the model to learn.

You can optionally add a "weight": 0 field to any message to exclude it from the loss computation. This is useful when you want the model to see context but only learn from specific responses.

The Training Pipeline

When you call client.fine_tuning.jobs.create(), Azure kicks off a pipeline:

Validation: Checks your JSONL for formatting errors, token limits, and minimum example counts (at least 10 examples required).
Queuing: Your job waits for GPU capacity.
Training: The model is fine-tuned using LoRA (more on this shortly). Azure automatically creates checkpoints.
Results: A results.csv file is generated with training and validation loss at each step.

Hyperparameters

You can customise the training with the hyperparameters argument:

Parameter	Default	What It Controls
`n_epochs`	Auto (based on dataset size)	Number of passes through the training data
`learning_rate_multiplier`	Auto	Scales the base learning rate. Higher means faster but riskier.
`batch_size`	Auto	Examples per gradient update. Larger is more stable but uses more memory.
`seed`	None	For reproducibility

job = client.fine_tuning.jobs.create(
    training_file=training_file.id,
    model="gpt-4o-mini-2024-07-18",
    hyperparameters={
        "n_epochs": 3,
        "learning_rate_multiplier": 1.8,
        "batch_size": 4,
    },
)

For most use cases, the defaults work well. Azure auto-selects values based on your dataset size. If you want to systematically search for optimal values, hyperparameter optimisation methods like Bayesian optimisation can help.

Pricing

Fine-tuning costs vary by model tier:

Tier	Training Cost	Hosting Cost	Best For
Standard	Higher per-token	Dedicated deployment	Production workloads
Global Standard	Moderate	Pay-per-use	Cost-effective production

Training costs are measured in tokens processed. A dataset of 1,000 examples at ~200 tokens each, trained for 3 epochs, processes about 600K tokens, typically a few dollars for smaller models like gpt-4o-mini.

Going Deeper

What Fine-Tuning Actually Does

A language model is a probability distribution over the next token. When you prompt GPT-4o with "The capital of France is", it assigns high probability to "Paris" and low probability to "pizza". These probabilities are determined by the model's weights, billions of numbers learned during pre-training.

Fine-tuning shifts these probability distributions. For our medical coding assistant, the base model might assign:

P("J20.9") = 0.001 (it's seen ICD codes, but rarely)
P("The patient has") = 0.15 (a more "natural" continuation)

After fine-tuning on hundreds of medical coding examples, the distribution shifts:

P("J20.9") = 0.85
P("The patient has") = 0.002

The training objective is cross-entropy loss on the assistant tokens, the same maximum likelihood estimation objective used in pre-training, just applied to a much smaller dataset. The model learns to maximise the probability of producing exactly the outputs in your training data.

The gradients that update the weights flow through the same backpropagation algorithm used in pre-training. The difference is scope: pre-training processes trillions of tokens across the entire internet; fine-tuning processes thousands of tokens from your specific domain.

LoRA: The Algorithm Under the Hood

Here's the problem with naive fine-tuning: GPT-4o has hundreds of billions of parameters. Updating all of them requires enormous GPU memory and risks catastrophically forgetting what the model already knows. This is where LoRA (Low-Rank Adaptation) comes in, and it's what Azure uses under the hood.

The key insight from Hu et al. (2021): when you fine-tune a large language model, the weight updates have low intrinsic rank. In plain English, the changes needed to adapt a model to a new task live in a much smaller subspace than the full parameter space.

Instead of updating a weight matrix $W \in \mathbb{R}^{d \times k}$ directly, LoRA decomposes the update into two smaller matrices:

Where:

$W$ is the original frozen weight matrix ( $d \times k$ )
$B \in \mathbb{R}^{d \times r}$ and $A \in \mathbb{R}^{r \times k}$ are the trainable low-rank matrices
$r \ll \min(d, k)$ is the rank, typically 8, 16, or 64

The parameter reduction is dramatic. Consider a weight matrix in a large transformer:

Original: $d = 4096, k = 4096$ → 16.8 million parameters
LoRA with $r = 16$ : $(4096 \times 16) + (16 \times 4096)$ = 131,072 parameters
Reduction: 99.2%

Across the entire model, LoRA typically trains only 0.1–1% of the original parameters. This means:

Memory: You can fine-tune on a single GPU instead of a cluster.
Speed: Fewer parameters to update means faster training.
Storage: Each fine-tuned version is just the small $B$ and $A$ matrices, megabytes instead of gigabytes.
No extra inference latency: At deployment, $BA$ is merged back into $W$ . The final model has exactly the same architecture and speed as the original.

The initialisation matters too: $B$ is initialised to zero and $A$ to random Gaussian, so $\Delta W = BA = 0$ at the start. Training begins from the exact pre-trained model, with no disruption.

The PEFT Family

LoRA belongs to a broader family called Parameter-Efficient Fine-Tuning (PEFT) methods. Here's how they compare:

Method	Trainable Params	Memory	Quality	Inference Overhead
Full fine-tuning	100%	Very high	Best	None
LoRA	0.1–1%	Low	Near-full	None (merged)
QLoRA	0.1–1%	Very low	Good	Slight (quantisation)
Prefix tuning	Under 0.1%	Very low	Moderate	Slight (extra tokens)
Adapters	1–5%	Low	Good	Slight (extra layers)

LoRA is the most popular because it hits the sweet spot: near-full fine-tuning quality with no inference overhead. QLoRA adds 4-bit quantisation of the base model, reducing memory further. You can fine-tune a 65B parameter model on a single 48GB GPU. Prefix tuning prepends learnable "virtual tokens" to the input, but quality degrades for complex tasks. Adapters insert small trainable layers between existing transformer blocks.

The Moat Spectrum

Not all AI customisation is equal. Here's the full spectrum, from least to most defensible:

Approach	Cost	Setup Time	Moat Strength	When to Use
Prompt engineering	Free	Minutes	None	Prototyping, one-off tasks
RAG (Retrieval-Augmented Generation)	$10s–100s/mo	Days	Weak (data can be copied)	Need current information, citations
Fine-tuning	$100s–1,000s	Days–weeks	Moderate (behaviour is learned)	Consistent formatting, domain tone, cost at scale
Pre-training	$10Ms–100Ms+	Months	Strong (architecture + data)	You're OpenAI, Google, or Meta

Prompt engineering is where most people stop. It works surprisingly well but offers zero competitive moat; anyone can copy your prompt. RAG adds your own data at inference time, which is powerful for knowledge-intensive tasks but the behaviour is still the base model's. Fine-tuning embeds behaviour into the weights. The model doesn't need to be told how to respond; it just does. Pre-training is building the engine from scratch, and unless you have a few hundred million dollars and a research lab, it's not your game.

When Fine-Tuning Beats Prompting

Fine-tuning wins over prompting when you need:

Consistent output formatting. JSON schemas, code conventions, structured reports. A fine-tuned model follows the format without lengthy system prompts.
Domain-specific behaviour. Medical coding, legal analysis, financial compliance. The model internalises domain norms.
Tone and style. Brand voice, technical writing style, conversational patterns.
Cost at scale. A fine-tuned model with a short prompt is cheaper per request than a base model with a 2,000-token system prompt.
Latency. Shorter prompts mean fewer input tokens to process.

When Fine-Tuning Loses to RAG

Fine-tuning embeds knowledge into weights, but weights are frozen after training. If the information changes frequently (stock prices, medical guidelines, product catalogues), RAG is the better choice. RAG retrieves current documents at inference time, so the model always has access to the latest information.

The best systems often combine both: fine-tune for behaviour (how to respond), RAG for knowledge (what to respond with).

Where This Comes From

LoRA: Low-Rank Adaptation of Large Language Models

LoRA was introduced by Hu et al. (2021) at Microsoft Research. The paper's central hypothesis is elegant:

"We hypothesize that the change in weights during model adaptation also has a low 'intrinsic rank,' which leads us to propose Low-Rank Adaptation (LoRA)."

The authors demonstrated that for GPT-3 175B, LoRA with rank 4 matched or exceeded full fine-tuning performance on multiple benchmarks while training only 0.01% of the parameters. They tested on natural language understanding (GLUE), natural language generation (E2E NLG), and instruction following. LoRA matched full fine-tuning across the board.

A key practical insight from the paper: LoRA is most effective when applied to the attention weight matrices ( $W_Q$ and $W_V$ ), rather than the feed-forward layers. This is because attention matrices control the model's "routing" of information (which tokens attend to which), and task-specific behaviour is largely about changing these routing patterns.

ULMFiT: The Transfer Learning Paradigm

Before LoRA, there was ULMFiT. Howard & Ruder (2018) established the now-standard paradigm: pre-train on a large corpus, then fine-tune on your task. Their key innovations, discriminative fine-tuning (different learning rates per layer) and gradual unfreezing, are the conceptual ancestors of LoRA's approach.

The Broader Lineage

The idea that pre-trained representations can be adapted to new tasks has a long history:

ImageNet transfer learning (2012–2014). Training on ImageNet, fine-tuning on medical images. Computer vision proved the concept.
ULMFiT (2018). Brought transfer learning to NLP. Demonstrated that language model pre-training produces universal features.
BERT (2018) and GPT (2018). Scaled the paradigm. Pre-train once, fine-tune for everything.
LoRA (2021). Made fine-tuning efficient enough for massive models. You don't need to update every parameter.

Each step reduced the barrier. LoRA's contribution is making fine-tuning feasible for models so large that full fine-tuning would require a cluster.

Try It Yourself

The format experiment. Take a task where you want structured output (e.g., JSON with specific fields). Compare: (a) a detailed system prompt describing the format, vs (b) a fine-tuned model trained on 50 examples of the correct format. Measure how often each produces valid output.
Data quality vs quantity. Create two training sets for the same task: 50 carefully curated, high-quality examples vs 500 noisy, auto-generated examples. Fine-tune on each. Quality almost always wins. This is the moat.
The moat test. Fine-tune a model on a specific domain task. Then try to replicate the same behaviour using only prompt engineering. How close can you get? Where does prompting fall short?
LoRA from scratch. Implement a toy LoRA layer in PyTorch. Freeze a pre-trained GPT-2 model, add $BA$ matrices to the attention layers, and fine-tune on a small text classification task. Compare the parameter count to full fine-tuning.

Interactive Tools

Explore Our Free Tools — Hands-on calculators and visualisers for statistics, machine learning, and quantitative finance

Maximum Likelihood Estimation from Scratch. Fine-tuning's loss function (cross-entropy) is maximum likelihood estimation. Understanding MLE gives you intuition for what the training loop is optimising.
Backpropagation and Neural Nets from First Principles. The gradient computation that makes both pre-training and fine-tuning work. LoRA reduces the number of parameters, but the gradients still flow through the same algorithm.
Hyperparameter Optimisation: Grid, Random, and Bayesian. Fine-tuning has its own hyperparameters (learning rate, epochs, LoRA rank). This post covers systematic approaches to tuning them.

Frequently Asked Questions

What is the difference between fine-tuning and prompt engineering?

Prompt engineering gives instructions to a base model at inference time, while fine-tuning embeds behaviour directly into the model's weights through additional training. Fine-tuning produces more consistent outputs without lengthy system prompts and can reduce per-request costs at scale. However, prompt engineering requires zero setup and is the right starting point for prototyping.

How much training data do I need for fine-tuning?

Azure OpenAI requires a minimum of 10 examples, but practical results typically need 50 to 500 high-quality examples depending on task complexity. Data quality matters far more than quantity: 50 carefully curated examples often outperform 500 noisy ones. Start small, evaluate, and add more data only if the model underperforms.

Does fine-tuning change the entire model?

No. Modern fine-tuning uses LoRA (Low-Rank Adaptation), which freezes the original model weights and trains only small low-rank matrices added to the attention layers. This typically updates only 0.1 to 1% of the original parameters, making fine-tuning feasible on modest hardware while preserving the base model's general capabilities.

Can I fine-tune open-source models instead of using Azure?

Yes. Open-source models like Llama and Mistral can be fine-tuned locally using libraries such as Hugging Face PEFT and QLoRA. The LoRA algorithm is the same regardless of platform. The trade-off is that you manage the infrastructure yourself, but you gain full control over the model and avoid ongoing API costs.

When should I use RAG instead of fine-tuning?

Use RAG (Retrieval-Augmented Generation) when the knowledge your model needs changes frequently, such as product catalogues, medical guidelines, or pricing data. Fine-tuning embeds knowledge into frozen weights, so it cannot adapt to new information without retraining. The best systems often combine both: fine-tune for consistent behaviour and formatting, then use RAG to inject up-to-date knowledge at inference time.

What is QLoRA and how does it differ from LoRA?

QLoRA combines LoRA with 4-bit quantisation of the base model, reducing memory requirements even further. With QLoRA, you can fine-tune a 65-billion parameter model on a single 48GB GPU. The trade-off is a slight quality reduction from quantisation and marginally higher inference latency compared to standard LoRA.

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.

DEV Community: Berkan Sesen

Value Iteration vs Q-Learning: Dynamic Programming Meets RL

Quick Win: Run Both Algorithms

What Just Happened?

Value Iteration: Planning with a Blueprint

Q-Learning: Learning by Doing

Why Both Reach the Same Answer

Going Deeper

Why Not 100%? The Stochasticity Tax

Convergence: 184 Sweeps vs 10,000 Episodes

The Model-Based vs Model-Free Tradeoff

Hyperparameter Sensitivity

When to Use Which

Deep Dive: The Papers

Bellman's Foundation

Watkins' Q-Learning

The DP-RL Connection

Further Reading

Try It Yourself

Interactive Tools

Related Posts

Frequently Asked Questions

What is the difference between value iteration and Q-learning?

What is the Bellman equation?

When should I use dynamic programming instead of Q-learning?

What is the difference between value iteration and policy iteration?

Does value iteration always converge?

Custom Likelihoods in PyMC: One-Inflated Beta Regression for Loan Repayment

Let's Build It

What Just Happened?

Two populations, one model

The piecewise log-likelihood

Hand-coding the Beta log-density

The model structure

Checking the fit

Posterior predictive check

Going Deeper

The mean-precision parameterisation

Why logistic links?

The expected value formula

Why pm.Potential and not pm.CustomDist?

Sampling considerations

When to use something else

Where This Comes From

Beta regression: Ferrari and Cribari-Neto (2004)

Inflated models: Ospina and Ferrari (2010)

From MLE to MCMC

Algorithm summary

Further reading

Interactive Tools

Related Posts

Frequently Asked Questions

When should I use a One-Inflated Beta model instead of logistic regression?

Why use pm.Potential instead of pm.CustomDist for the likelihood?

What is the mean-precision parameterisation of the Beta distribution?

How do I check whether the OIB model fits my data well?

Can this model handle spikes at both 0 and 1?

Bayesian Survival Analysis with PyMC: Modelling Customer Churn

Let's Build It

What Just Happened?

Right-Censoring: Learning from Incomplete Data

Why Gumbel? The Weibull-Gumbel Connection

pm.Potential: Telling PyMC About Partial Information

MCMC Diagnostics

Survival Curves from the Posterior

Going Deeper

Covariates in the Scale Too

Weibull vs Log-Logistic: Which Tail Shape?

The Cox Proportional Hazards Alternative

When NOT to Use Bayesian AFT

Where This Comes From

Cox (1972): Proportional Hazards

Buckley and James (1979): Accelerated Failure Time

Wei (1992): AFT as an Alternative

Handling Censoring in PyMC

Further Reading

Interactive Tools

Related Posts

Frequently Asked Questions

What is right-censoring and why does it matter?

`pm.Potential`: Telling PyMC About Partial Information