Andreas Bogossian https://andreasbogossian.com/blog.html Notes on challenging technical things quarto-99.9.9 Sat, 25 Apr 2026 21:00:00 GMT The Kruskal-Wallis test, derived from first principles Andreas Bogossian https://andreasbogossian.com/posts/Kruskal-Wallis-test/ Let’s say that we want to test the effects of pain management protocols. We have 45 post-surgical patients randomized to three groups: opioid-only, NSAID + nerve block and multimodal. 48 hours after their operations, each patient marks their pain on a 0–10 visual analog scale (VAS). The question that the clinic is interested in is whether the three different protocols produce different pain levels or not.

ANOVA [1] is a popular test for testing whether group means are equal. The problem with ANOVA in the pain scale example is that the pain scores are typically right-skewed (clustering around 2-4) and with small sample sizes, the within-group variance can be inflated by some observations being at the right tail (8-10). ANOVA computes an F-statistic as a ratio of between-group variance to within-group variance. If the within group variance is inflated, ANOVA needs a larger between-group variance to detect a difference.

The Kruskal-Wallis test solves this problem. The KW test replaces raw scores with their ranks in the pooled dataset (see Figure 2) and tests whether the ranks cluster by treatment. Under the null hypothesis, the ranks should equally be distributed to the groups.

Symbol Meaning
Number of groups
Group index,
Number of observations in group
Total sample size
The -th observation in group
Rank of among all pooled observations
Rank sum of group :
Mean rank of group :
Grand mean rank:
Kruskal-Wallis test statistic
ANOVA F-statistic
Chi-square distribution with degrees of freedom

The pain trial

This post will use a made up example: the data will have three groups of 15 patients with NRS scores at 48 hours post surgery. The Figure 1 shows both the theoretical (Gamma) distributions where the samples are drawn from and the actual samples. Then Table 1 summarizes the data.

Pain trial data 
import numpy as np
import matplotlib.pyplot as plt
from scipy import stats

# VAS pain scores at 48 h post-op (0 = no pain, 10 = worst imaginable)
# Gamma(shape=0.5) → heavy right skew; offset shifts group medians apart
rng = np.random.default_rng(99)
pain_a = np.clip(np.round(rng.gamma(0.5, 3.0, 15) + 1.0, 2), 0, 10)  # multimodal:          mean ≈ 2.6
pain_b = np.clip(np.round(rng.gamma(0.5, 3.0, 15) + 1.8, 2), 0, 10)  # NSAID + nerve block: mean ≈ 3.4
pain_c = np.clip(np.round(rng.gamma(0.5, 3.0, 15) + 2.6, 2), 0, 10)  # opioid-only:         mean ≈ 3.9

labels = ["Multimodal", "NSAID + nerve block", "Opioid-only"]
colors = ["mediumseagreen", "darkorange", "steelblue"]
groups = [pain_a, pain_b, pain_c]

plt.rcParams.update(
    {
        "text.color": "white",
        "axes.labelcolor": "white",
        "xtick.color": "white",
        "ytick.color": "white",
        "axes.edgecolor": "white",
    }
)

fig, axes = plt.subplots(1, 3, figsize=(10, 3.2), sharey=True, sharex=True)
fig.patch.set_alpha(0)

for ax, g, color, label in zip(axes, groups, colors, labels):
    ax.hist(
        g,
        bins=np.linspace(0, 10, 12),
        color=color,
        alpha=0.8,
        edgecolor="white",
        linewidth=0.8,
    )
    ax.axvline(np.median(g), color="white", linewidth=2, linestyle="--", alpha=0.9)
    ax.set_title(f"{label}\nmedian = {np.median(g):.1f}", fontsize=10, color="white")
    ax.set_xlabel("Pain score (0–10)")
    ax.set_xlim(-0.5, 10.5)
    ax.patch.set_alpha(0)

axes[0].set_ylabel("Count")
plt.tight_layout()
plt.show()
Figure 1: VAS pain scores at 48 h post-surgery for three protocols. Dashed line shows the group median.
Summary statistics 
from IPython.display import HTML

rows = "".join(
    f"<tr><td style='padding:5px 14px;'>{lbl}</td>"
    f"<td style='padding:5px 14px; text-align:right;'>{len(g)}</td>"
    f"<td style='padding:5px 14px; text-align:right;'>{g.mean():.2f}</td>"
    f"<td style='padding:5px 14px; text-align:right;'>{np.median(g):.1f}</td>"
    f"<td style='padding:5px 14px; text-align:right;'>{g.std(ddof=1):.2f}</td></tr>"
    for lbl, g in zip(labels, groups)
)
th = "padding:6px 14px; text-align:right; border-bottom:1px solid #555;"
th_l = "padding:6px 14px; text-align:left; border-bottom:1px solid #555;"
display(HTML(f"""
<table style="border-collapse:collapse; font-size:0.9em;">
  <thead><tr>
    <th style="{th_l}">Protocol</th>
    <th style="{th}">n</th>
    <th style="{th}">Mean</th>
    <th style="{th}">Median</th>
    <th style="{th}">SD</th>
  </tr></thead>
  <tbody>{rows}</tbody>
</table>
"""))
Table 1: Summary statistics for the three pain management protocols.
Protocol n Mean Median SD
Multimodal 15 2.57 1.8 2.21
NSAID + nerve block 15 3.21 3.3 1.37
Opioid-only 15 3.93 3.2 1.92

The key insight: order is enough

For deducing whether one protocol produces less pain than another, we don’t need to know how much less. We only need to know who ranks lower. Thus we use a ranking transformation. The animation in Figure 2 shows how the data points are transformed into ranks.

Figure 2: Each dot is one of the 45 patients, coloured by protocol. Drag to 1: each pain score is replaced by its rank in the pooled dataset, spreading the dots uniformly. The pooled row fades in to show all 45 ranks in sequence.

Deriving the statistic

Setting up the problem

Abstracting from the medical example, let’s say that we have groups with the number of observations in each group being . Then the total number of observations is . After combining the observations into one pool and ranking them (as shown in Figure 2), let be the rank of the -th observation in group .

The goal is to test whether the groups share the same distribution with each other. Transforming this into a null hypothesis becomes

So the next question to answer is: what does a large deviation from the random distribution look like and how can we measure the deviation?

Expected rank mean under

Under , every rank is equally likely to fall into any group. Then the expected value of a single rank drawn from the uniform set is

The expected mean rank in all groups is , regardless of the individual group sizes. The only thing that matters is the total sample size. If is true, all groups’ mean ranks should hover around .

Building the test statistic

A natural way to measure deviation from the null is the squared distance of each group’s mean rank from the expected grand mean:

Here is the size of group to weight the penalty of bigger groups more because we have more information of them. This structure is extremely similar to ANOVA that is defined as the between group sum of squared deviations from each group mean from the global mean:

The scaling factor

To turn from Equation 2 into a statistic with a known null distribution, we need to normalize it by the variance of the ranks. With Equation 1, the variance for a discrete uniform distribution on is:

By analogy with the ANOVA F-statistic, we normalise by , where the total rank sum-of-squares for the complete set of ranks is :

The factor scales so that its expected value under equals , which is exactly the degrees-of-freedom parameter of a chi-square distribution. For large , under .

Back to the pain trial

Now that we’ve shown the mathematical difference between ANOVA and the Kruskal-Wallis test, we can compute the p-values using both of the tests and compare them.

import numpy as np
from scipy import stats

N = sum(len(g) for g in groups)
K = len(groups)

all_scores = np.concatenate(groups)
ranks = stats.rankdata(all_scores)

idx = 0
rank_groups = []
for g in groups:
    rank_groups.append(ranks[idx : idx + len(g)])
    idx += len(g)

grand_mean_rank = (N + 1) / 2
H, p_kw = stats.kruskal(*groups)
F, p_anova = stats.f_oneway(*groups)
Table rendering 
from IPython.display import HTML

th = "padding:6px 14px; text-align:right; border-bottom:1px solid #555;"
th_l = "padding:6px 14px; text-align:left; border-bottom:1px solid #555;"
td = "padding:5px 14px; text-align:right;"
td_l = "padding:5px 14px;"

rank_rows = "".join(
    f"<tr><td style='{td_l}'>{lbl}</td>"
    f"<td style='{td}'>{rg.sum():.0f}</td>"
    f"<td style='{td}'>{rg.mean():.2f}</td>"
    f"<td style='{td}'>{rg.mean() - grand_mean_rank:+.2f}</td></tr>"
    for lbl, rg in zip(labels, rank_groups)
)

test_rows = (
    f"<tr><td style='{td_l}'>Kruskal-Wallis</td>"
    f"<td style='{td}'>H = {H:.3f}</td>"
    f"<td style='{td}'>{K - 1}</td>"
    f"<td style='{td}'>{p_kw:.4f}</td></tr>"
    f"<tr><td style='{td_l}'>ANOVA</td>"
    f"<td style='{td}'>F = {F:.3f}</td>"
    f"<td style='{td}'>{K - 1}</td>"
    f"<td style='{td}'>{p_anova:.4f}</td></tr>"
)

display(HTML(f"""
<p style="font-size:0.85em;color:#aaa;margin-bottom:4px;">
  Grand mean rank (null expectation): {grand_mean_rank:.1f}
</p>
<table style="border-collapse:collapse;font-size:0.9em;margin-bottom:1.2em;">
  <thead><tr>
    <th style="{th_l}">Protocol</th>
    <th style="{th}">Rank sum</th>
    <th style="{th}">Mean rank</th>
    <th style="{th}">Deviation</th>
  </tr></thead>
  <tbody>{rank_rows}</tbody>
</table>
<table style="border-collapse:collapse;font-size:0.9em;">
  <thead><tr>
    <th style="{th_l}">Test</th>
    <th style="{th}">Statistic</th>
    <th style="{th}">df</th>
    <th style="{th}">p-value</th>
  </tr></thead>
  <tbody>{test_rows}</tbody>
</table>
"""))

Grand mean rank (null expectation): 23.0

Protocol Rank sum Mean rank Deviation
Multimodal 212 14.17 -8.83
NSAID + nerve block 378 25.20 +2.20
Opioid-only 444 29.63 +6.63
Test Statistic df p-value
Kruskal-Wallis H = 11.034 2 0.0040
ANOVA F = 1.976 2 0.1513

The multimodal group’s mean rank is well below 23.0 (the grand mean rank), the opioid-only group’s mean rank is well above it. The NSAID + nerve block is quite close to the mean rank. The -statistic is able to quantify the spread. The chi-square p-value tells us whether a deviation this large would occur by chance under the null. With , we reject. The KW-test rejects the null, ANOVA doesn’t. This shows that the KW-test is able to quantify the differences between distributions when they are heavily skewed better than ANOVA.

NoteWhat KW doesn’t tell

A statistically significant KW result only says that something differs, not which pairs differ. To investigate which pairs differ, the Dunn test [2] fits this problem well.

References

[1]
Wikipedia contributors, “Analysis of variance.” https://en.wikipedia.org/wiki/Analysis_of_variance, 2025.
[2]
Z. Bobbitt, “Dunn’s test for multiple comparisons.” https://www.statology.org/dunns-test/, 2020.

Reuse

Citation

BibTeX citation:
@online{bogossian2026,
  author = {Bogossian, Andreas},
  title = {The {Kruskal-Wallis} Test, Derived from First Principles},
  date = {2026-04-26},
  url = {https://andreasbogossian.com/posts/Kruskal-Wallis-test/},
  langid = {en}
}
For attribution, please cite this work as:
A. Bogossian, “The Kruskal-Wallis test, derived from first principles.” [Online]. Available: https://andreasbogossian.com/posts/Kruskal-Wallis-test/
]]> statistics hypothesis testing https://andreasbogossian.com/posts/Kruskal-Wallis-test/ Sat, 25 Apr 2026 21:00:00 GMT The EM-algorithm Andreas Bogossian https://andreasbogossian.com/posts/em-algorithm/ Let’s start from a practical example. Let’s say that a teacher has their 220 students choose between a difficult exam or an easier exam. When the exam papers are collected, the teacher accidentally messes up the papers and doesn’t know which scores are from the difficult exam and which are from the easer one. Thus the teacher is left with a pile of 220 unlabelled exam scores.

Now what would be the best way to distinguish which exam result comes from which exam? One solution would be to just draw a line where e.g. everyone under 67 points on the exam took the more difficult exam and everyone over took the easier one. The problem with this approach is that the point distributions overlap. This means that some students achieving great points on the more difficult exam would be incorrectly labelled as the worst performing people on the easier exam.

A better approach is to fit who Gaussian models to soft label the exam scores. This way, every exam result would have a probability that would tell the teacher the chances of the point coming from each category. Now in order to soft label the exam scores, we need to know the parameters for the Gaussian distributions. But to know what parameters fit the models to the data the best, we need the information about which group each data point belongs to.

The problem is that this is a “chicken-and-egg” problem, where in order to acquire one, you need the other. When you don’t know either, you can’t determine either. This deadlock can be broken by alternating between guessing the labels and updating the distributions until distributions converge. This method is called the EM-algorithm.

Symbol Name Description
Observed data The full set of exam scores
Single observation The score of student
Latent variables The hidden group labels for all scores
Single latent variable The hidden group label for student
Parameters All model parameters
Parameters at iteration The current parameter estimate during EM
Sample size The total number of exam scores
Number of components The number of Gaussian distributions in the mixture
Component index Indexes a single Gaussian component,
Mixing weight How common group is; all sum to 1
Mean Centre of the -th Gaussian
Variance Spread of the -th Gaussian
Gaussian density PDF of a Normal distribution evaluated at
Log-likelihood Log probability of the observed data under
Complete-data log-likelihood Log probability if the hidden labels were known
Q-function Expected complete-data log-likelihood under the current posterior of
Responsibility Posterior probability that score came from component
Indicator function 1 if the condition inside is true, 0 otherwise
Convergence threshold Minimum log-likelihood change required to keep iterating

Theory

The model: Gaussian Mixture Model

We model the data as coming from a Gaussian Mixture Model (GMM). A GMM is a weighted sum of Normal distributions. The marginal likelihood (the probability of the data given the parameters) is:

where

  • is mixing weight. It asks how common is group ? Must sum to 1.
  • are the mean and variance of the -th Gaussian.
  • is the latent group assignment for point , which we don’t know.

Now in order to soft label the exam scores, we want to find that maximises the log-likelihood of the data. The log likelihood can be written as:

But there is a problem: The log of a sum has no tidy closed form derivative and this is why plain maximum likelihood fails [1].

The EM fix: complete-data log-likelihood

The trick the EM algorithm uses is that it imagines we did observe the hidden labels . So If we know which group each score came from, the joint probability of a single observation would be:

Over all points, the complete-data log-likelihood is:

Since each takes exactly one value in , we can rewrite this using an indicator that makes the sum over k explicit

In this notation is 1 when point is truly from group and 0 otherwise. Now the reason we write the likelihood in this for is because the sum of logs is tractable and can be used for optimization. But there is a catch: we still don’t know , so we can’t evaluate directly. Instead, we take the expectation of over the posterior of given the data and current parameters :

Expanding the expectation, the hard indicators get replaced by soft posterior probabilities :

This form is a weighted sum of logs that is differentiable and maximisable in closed form. We can derive the meaning of the two steps from this form:

  • E-step: compute (estimate) the weights
  • M-step: maximise over treating as fixed

So in summary this means that we solve the chicken-and-egg problem by estimating the responsibilities and then estimating the parameters based on the estimate of the responsibility. This way we always have one of the two unknowns locked and can create an estimate from the other unknown.

E-step: compute responsibilities

The weights are the posterior probability that the point came from the component , given the current parameters. By Bayes’ theorem, prior belief about group membership gets updated by how well each Gaussian explains the observed score:

In the fraction, the numerator is the prior belief that point is from group (), multiplied by how well group ’s Gaussian explains the score (). The denominator then normalizes across all groups so the responsibilities for each point sum up to 1. The result is then a soft assignment so every point gets a fractional membership in every group rather than a binary 0/1 label. After this step we have an matrix of responsibilities.

M-step: update parameters

With responsibilities fixed, we maximize over by differentiating with respect to each parameter, setting to zero and solving. The resulting updates are weighted statistics where each data point contributes to component in proportion to [1, Ch. 9].

Here

  • is the responsibility-weighted mean of the data
  • is the responsibility-weighted variance around that mean
  • is the average responsibility — the fraction of the data “owned” by component

Convergence

Now to solve the optimal parameters, we iterate alternating the E and M steps until the change in the log-likelihood falls below a threshold . A thing to keep in mind with the EM algorithm is that it is guaranteed to reach a local maximum (or a saddle point) and not necessary the global one [1, Ch. 9]. Also initialization matters because different starting points can lead to different local optima.

Code demo

In this code demo, we’ll use the example from the introduction. We’ll create a data set from two exams that are sampled from two distinct Gaussian distribution. The goal is to soft label the exam scores so that students can get their correct grades based on exam difficulty they chose.

Setup and data generation

import numpy as np
import json
from scipy.stats import norm

rng = np.random.default_rng(42)

# Generate synthetic "exam score" data from two groups:
#   Group 1 (hard paper): mean=45, sd=8
#   Group 2 (easy paper): mean=70, sd=6
n1, n2 = 120, 100
group1 = rng.normal(loc=45, scale=8, size=n1)
group2 = rng.normal(loc=70, scale=6, size=n2)
true_labels = np.array([0] * n1 + [1] * n2)
idx = rng.permutation(n1 + n2)
X = np.concatenate([group1, group2])[idx]
true_labels = true_labels[idx]
n = len(X)
Code 
import matplotlib.pyplot as plt
from scipy.stats import norm as sp_norm

fig, ax = plt.subplots(figsize=(8, 3.5))
ax.hist(X, bins=30, density=True, color="#cbd5e1", edgecolor="white", label="Observed scores")

xs_plot = np.linspace(10, 100, 300)
ax.plot(xs_plot, sp_norm.pdf(xs_plot, 45, 8),  color="steelblue",  lw=2, label="True group 1 (μ=45, σ=8)")
ax.plot(xs_plot, sp_norm.pdf(xs_plot, 70, 6),  color="darkorange", lw=2, label="True group 2 (μ=70, σ=6)")

ax.set_xlabel("Score")
ax.set_ylabel("Density")
ax.legend()
plt.tight_layout()
plt.show()

There are 220 exam scores in the dataset sampled from two overlapping Gaussian distributions. After sampling the exam scores, we shuffle the scores so that the algorithm can’t distinguish what scores come from what distribution. We include the original labels only for testing purposes to visualize how well the EM algorithm is able to find the correct labels. The EM algorithm follows roughly the following steps:

  1. Start with a rough guess at where the two groups are
  2. For each data point, calculate the probability it belongs to group 1 vs group 2
  3. Use those probabilities as soft weights to refit the distributions. Points clearly in one group pull hard on that group’s mean, uncertain points contribute a little to both
  4. After this we have slightly better distributions, which give slightly better probabilities next time
  5. Repeat until both stabilize

EM implementation

def em_gmm(X, K=2, n_iter=30, seed=0):
    rng_init = np.random.default_rng(seed)
    n = len(X)

    # --- Initialise parameters randomly ---
    mu    = rng_init.choice(X, K, replace=False)   # pick K random data points as starting means
    sigma = np.full(K, np.std(X))                  # same initial spread for all components
    pi    = np.ones(K) / K                          # equal mixing weights

    history = []   # store every snapshot for the OJS step-through

    for t in range(n_iter):

        # --- E-step: responsibilities ---
        # r[i, k] = p(z_i = k | x_i, theta^(t))
        r = np.array([
            pi[k] * norm.pdf(X, mu[k], sigma[k])
            for k in range(K)
        ]).T                              # shape: (n, K)
        r /= r.sum(axis=1, keepdims=True) # normalise rows

        # --- M-step: update parameters ---
        Nk    = r.sum(axis=0)
        mu    = (r * X[:, None]).sum(axis=0) / Nk
        sigma = np.sqrt((r * (X[:, None] - mu) ** 2).sum(axis=0) / Nk)
        pi    = Nk / n

        # --- Log-likelihood ---
        ll = np.log(
            sum(pi[k] * norm.pdf(X, mu[k], sigma[k]) for k in range(K))
        ).sum()

        history.append({
            "iter":  t + 1,
            "mu":    mu.tolist(),
            "sigma": sigma.tolist(),
            "pi":    pi.tolist(),
            "ll":    round(ll, 3),
            "r":     r.tolist(),   # full responsibility matrix. Used for rug colouring
        })

    return history


history = em_gmm(X, K=2, n_iter=30)

# Bridge to OJS
ojs_define(
    history     = history,
    X_data      = X.tolist(),
    true_labels = true_labels.tolist(),
    true_mu     = [45.0, 70.0],
    true_sigma  = [8.0,  6.0],
)

Interactive visualizer

The slider steps through the iterations of the EM algorithm. The top plot shows how the two Gaussians converge towards the most likely parameter values. The bottom plot shows the log-likelihood.

The two curves are the fitted Gaussians and the rug marks show each data point coloured by its soft assignment (blue = component 1, orange = component 2, yellow = uncertain). Grey dashed lines are the true means.

Log-likelihood across iterations. It can only ever increase.

References

[1]
C. M. Bishop, Pattern recognition and machine learning. Springer, 2006.

Reuse

Citation

BibTeX citation:
@online{bogossian2026,
  author = {Bogossian, Andreas},
  title = {The {EM-algorithm}},
  date = {2026-04-14},
  url = {https://andreasbogossian.com/posts/em-algorithm/},
  langid = {en}
}
For attribution, please cite this work as:
A. Bogossian, “The EM-algorithm.” [Online]. Available: https://andreasbogossian.com/posts/em-algorithm/
]]> probabilistic ML algorithms variational inference https://andreasbogossian.com/posts/em-algorithm/ Mon, 13 Apr 2026 21:00:00 GMT Variational Inference & the ELBO Andreas Bogossian https://andreasbogossian.com/posts/evidence-lower-bound/
Symbol Meaning
Observed data
Latent variables
True (intractable) posterior
Likelihood
Prior distribution
Marginal likelihood (evidence)
Variational approximation to the posterior
Generative model parameters
Variational parameters
Evidence Lower Bound (ELBO)
KL divergence from to
Expectation under distribution
Mean and std. dev. of the variational distribution
Mean and std. dev. of the true posterior
Likelihood standard deviation
Gaussian distribution with mean and variance

Let’s say we’ve observed the data and we want to know the conditional probability of the latent variables . The conditional probability distribution can be then formulated with the Bayes’ rule as

Now the problem is that the denominator that acts as a regularizer

is often intractable [1]. When something is intractable, it means that it is theoretically solvable but can’t be solved in a practical amount of time or with practical computational resources.

Variational Inference: Approximating the Posterior

With variational inference, we introduce a tractable approximation and find the parameters that approximate the intractable posterior as well as possible. The closeness is measured by KL-divergence. Even though KL-divergence is a good measure for determining the closeness of two distributions, it requires knowing the true posterior, which we don’t know because we don’t know the denominator which we’re trying to solve, but we need the posterior to solve the denominator. Using KL-divergence alone will not solve the approximation problem so we maximize a quantity called Evidence Lower BOund (ELBO).

Derivation of the ELBO [2]

Now our goal is to transform the intractable integral into a tractable approximation . Instead of approximating we approximate . So we start from

Now by multiplying by we get

Using the property

Next, we can apply the chain rule of probability: :

Multiplying by one, by introducing the new approximate tractable distribution :

With KL-divergence is defines as :

Since KL divergence is always non-negative, the ELBO is a lower bound on the log evidence:

The notation is used for the ELBO, so I’ll be also using it to simplify the math notation.

What is the ELBO Optimising?

Our original goal was to find a tractable distribution that approximates . A novel solution to this is to maximize the ELBO. With being fixed, maximizing the ELBO minimizes the KL-divergence.

Expanding the ELBO

The ELBO can be further decomposed by factoring using the chain rule of probability and splitting the log. This decomposition is a necessary tool for variational autoencoders (VAEs).

Here the reconstruction term measures how well the model recovers from samples of the latent space and the regularisation term penalises for deviating from the prior .

A Simple Example

Model:

True posterior (conjugate, so tractable here):

Variational family: .

We optimise by gradient ascent on the ELBO and compare with the truth.

np.random.seed(42)
x_obs   = 1.8
sigma_l = 1.2
mu_post   = x_obs / (sigma_l**2 + 1)
sig2_post = sigma_l**2 / (sigma_l**2 + 1)

def elbo_gaussian(mu_q, log_sq):
    sq = np.exp(log_sq)
    sq2 = sq**2
    recon = -0.5 * np.log(2 * np.pi * sigma_l**2) \
            - 0.5 * (sq2 + (mu_q - x_obs)**2) / sigma_l**2
    kl = 0.5 * (sq2 + mu_q**2 - 1 - 2*log_sq)
    return recon - kl

mu_q, log_sq = 0.0, 0.0
lr = 0.08
history = []
for _ in range(120):
    sq = np.exp(log_sq); sq2 = sq**2
    grad_mu  = (x_obs - mu_q)/sigma_l**2 - mu_q
    grad_lsq = -sq2/sigma_l**2 + 1 - sq2
    mu_q   += lr * grad_mu
    log_sq += lr * grad_lsq
    history.append({
        "iter":   len(history) + 1,
        "mu_q":   round(mu_q, 4),
        "sq":     round(np.exp(log_sq), 4),
        "elbo":   round(elbo_gaussian(mu_q, log_sq), 4),
    })

ojs_define(
    history   = history,
    mu_post   = mu_post,
    sig_post  = float(np.sqrt(sig2_post)),
)
(a)
(b)
(c)
(d)
(e)
(f)
Figure 1: VI convergence: the variational q approaches the true posterior. Drag the slider to step through iterations. Dashed blue line is the true posterior; solid orange line is the current q.
(a)
(b)
Figure 2: ELBO rises monotonically during VI optimisation. 
# True log evidence via numerical integration
z_grid = np.linspace(-10, 10, 10000)
dz = z_grid[1] - z_grid[0]
log_joint = stats.norm.logpdf(z_grid, x_obs, sigma_l) + stats.norm.logpdf(z_grid, 0, 1)
log_evidence_numerical = logsumexp(log_joint) + np.log(dz)

print(f"True log p(x) (numerical): {log_evidence_numerical:.4f}")
print(f"Final ELBO:                {history[-1]['elbo']:.4f}")
print(f"Difference:                {abs(log_evidence_numerical - history[-1]['elbo']):.4f}")
True log p(x) (numerical): -2.0289
Final ELBO:                -2.0289
Difference:                0.0000

As can be seen in Figure 1, the approximation approaches the true distribution. This is verified numerically in the code block above. The code block shows that the KL divergence has converged to zero meaning that the distributions are the same up until numerical precision.

Summary

In the beginning, we motivated the need for variational inference to solve intractable posteriors. Then we derived ELBO and showed that maximizing it is equivalent to minimizing the KL divergence between the approximation and the true posterior. Then finally, we used the theory on a simple Gaussian example confirming that we were able to approximate the posterior well.

References

[1]
C. M. Bishop, Pattern recognition and machine learning. Springer, 2006.
[2]
U. Jamil, “Variational Autoencoder - Model, ELBO, loss function and maths explained easily!” https://www.youtube.com/watch?v=iwEzwTTalbg, Jun. 2023.

Reuse

Citation

BibTeX citation:
@online{bogossian2026,
  author = {Bogossian, Andreas},
  title = {Variational {Inference} \& the {ELBO}},
  date = {2026-04-02},
  url = {https://andreasbogossian.com/posts/evidence-lower-bound/},
  langid = {en}
}
For attribution, please cite this work as:
A. Bogossian, “Variational Inference & the ELBO.” [Online]. Available: https://andreasbogossian.com/posts/evidence-lower-bound/
]]> variational inference probabilistic ML deep learning https://andreasbogossian.com/posts/evidence-lower-bound/ Wed, 01 Apr 2026 21:00:00 GMT A* Pathfinding Algorithm Andreas Bogossian https://andreasbogossian.com/posts/a-star/ How does Google Maps find the best route in under a second, across millions of roads? Stripping this question into its basic elements, the question becomes, how to find the shortest path in a graph from the start position to the end position. If one was to approach this problem by trying every single path, the number of paths grow as , where is the branching factor (average number of branches from each node) and is the depth. The brute force method scales extremely poorly so another solution must be used.

Dijkstra’s algorithm improves on the brute force method slightly. Dijkstra always expands the cheapest node so far. The problem with Dijkstra is that it expands the search in every direction including roads that lead away from the destination. Dijkstra can be beaten with algorithms with sense of direction.

The idea behind the A* algorithm

A* assigns every node a single score:

where

  • : the exact cost from start to the current node . This is known.
  • : the estimated cost from the current node to the goal. This is a heuristic (educated guess).

After A* computes the score of each frontier node, the algorithm will always expand the node with the lowest . Minimizing both the cost of getting to the current node and the estimate of arriving at the goal makes the algorithm seek the shortest path from start point to the end point.

For a more thorough study of A*, please refer to [1].

Implementation

Heuristic

A common heuristic is the Manhattan distance:

The geometric interpretation of this is amount of grid cells that need to be navigated to go from one point to another without moving diagonally.

def heuristic(a: tuple[int, int], b: tuple[int, int]) -> int:
    return abs(a[0] - b[0]) + abs(a[1] - b[1])

Main loop

The open set is a min-heap keyed on . At each iteration the node with the lowest is popped. Then check if the goal has been reached. Then the third step is to expand the neighbours:

import heapq

def astar(
    grid: list[list[int]],
    start: tuple[int, int],
    goal: tuple[int, int]
) -> dict[tuple[int, int], tuple[int, int]]:
    rows, cols = len(grid), len(grid[0])

    open_set: list[tuple[int, tuple[int, int]]] = []
    heapq.heappush(open_set, (0, start))

    came_from: dict[tuple[int, int], tuple[int, int]] = {}
    g_score: dict[tuple[int, int], int] = {start: 0}

    while open_set:
        _, current = heapq.heappop(open_set)

        if current == goal:
            break

        r, c = current
        for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
            neighbour: tuple[int, int] = (r + dr, c + dc)
            nr, nc = neighbour

            if not (0 <= nr < rows and 0 <= nc < cols):
                continue
            if grid[nr][nc] == 1:
                continue

            tentative_g = g_score[current] + 1

            if tentative_g < g_score.get(neighbour, float('inf')):
                came_from[neighbour] = current
                g_score[neighbour] = tentative_g
                f = tentative_g + heuristic(neighbour, goal)
                heapq.heappush(open_set, (f, neighbour))

    return came_from

Path reconstruction

def reconstruct_path(
    came_from: dict[tuple[int, int], tuple[int, int]],
    start: tuple[int, int],
    goal: tuple[int, int]
) -> list[tuple[int, int]]:
    path: list[tuple[int, int]] = []
    node = goal
    while node != start:
        path.append(node)
        node = came_from[node]
    path.append(start)
    path.reverse()
    return path

Running the algorithm on a 2D grid

grid = [
    [0, 0, 0, 1, 0],
    [1, 1, 0, 1, 0],
    [0, 0, 0, 0, 0],
    [0, 1, 1, 1, 0],
    [0, 0, 0, 0, 0],
]
start, goal = (0, 0), (4, 4)

came_from = astar(grid, start, goal)
path = reconstruct_path(came_from, start, goal)
print(path)
[(0, 0), (0, 1), (0, 2), (1, 2), (2, 2), (2, 3), (2, 4), (3, 4), (4, 4)]

Weighted A*

A useful, greedier A* variant is weighted A*. The variant inflates the heuristic by a weight :

The benefit of using the weight is that fewer nodes are expanded because the search prioritizes nodes that are close to the end. The trade-off is that the path found is at most times longer than optimal.

Demo

Start with Dijkstra and then swap to A* with the Manhattan heuristic. Also, test the weight slider to make the algorithm faster with the expense of accuracy.

  • Blue-green cells — open set: discovered, not yet expanded.
  • Orange cell — node currently being expanded.
  • Dark green cells — closed set: fully expanded.
  • Yellow path — the solution once the algorithm reaches the end node.
  • Drag S or E to reposition start and goal. Paint walls by clicking and dragging on empty cells.

Discussion

A* shines in applications where a cost function and an admissible heuristic (a heuristic that never overestimates the distance to the end) can be defined. Examples of these types of application are GPS routing, NPC path finding in games and network routing.

There is one problem with A*: it doesn’t factor walls into the heuristic when using Manhattan distance. For solving this problem, algorithms like Jump Point Search (JPS) are used.

References

[1]
S. Russell and P. Norvig, Artificial intelligence: A modern approach, 4th ed. Pearson Education Limited, 2022.

Reuse

Citation

BibTeX citation:
@online{bogossian2026,
  author = {Bogossian, Andreas},
  title = {A* {Pathfinding} {Algorithm}},
  date = {2026-03-19},
  url = {https://andreasbogossian.com/posts/a-star/},
  langid = {en}
}
For attribution, please cite this work as:
A. Bogossian, “A* Pathfinding Algorithm.” [Online]. Available: https://andreasbogossian.com/posts/a-star/
]]> algorithms search graph theory https://andreasbogossian.com/posts/a-star/ Wed, 18 Mar 2026 22:00:00 GMT KL Divergence Andreas Bogossian https://andreasbogossian.com/posts/KL-divergence/ When you fine-tune a large language model, how do you measure whether its output distribution is actually moving closer to the desired distribution? This question is a central to modern alignment techniques like RLHF [1], [2].

Large language models like GPTs output a probability distribution over tokens at each step. They assign a likelihood for each possible next word and then append the most likely next word to the output. During training the model outputs a predicted distribution of the most likely next words. For ensuring correct output, the distribution needs to be compared to the true distribution . This is exactly what KL-divergence does. KL-divergence takes the original distribution and compares it to a new generated distribution and outputs a value that tells how much the two distributions differ. The closer the value is to zero, the more similar the distributions are.

Theory

The Kullback-Leibler divergence from Q to P is defined as:

where and are the probabilities assigned by distributions and to outcome , and the sum is taken over all outcomes in the support of .

KL divergence can be interpreted intuitively using the concept of entropy form information theory.

The entropy of is

This is the average number of bits needed to encode samples from P using an optimal code for [3].

If instead if a code optimized for was used to encode samples that came form , the average code length is the cross-entropy:

Now combining entropy and cross entropy together, KL divergence is the extra cost you pay for using a code for to encode samples from :

Key properties

  • Asymmetric:

  • Always non-negative:

  • Equal to zero only when the two distributions are identical.

So intuitively, the KL-divergence asks how supriced you would be about the current distribution of the model given you know what the training data looks like. The more suprised you would be, the further the model is from the truth and thus the KL-divergence would be bigger.

Code demo

Now imagine a small vocabulary of five tokens: ["cat", "sat", "on", "the", "mat"]. We have an imaginary LLM that has a predicted distribution for the vocabulary and we also know the true distribution from the training data for the vocabulary.

import numpy as np

vocab = ["cat", "sat", "on", "the", "mat"]

# True distribution (training data)
P = np.array([0.4, 0.3, 0.1, 0.15, 0.05])

# Model's predicted distribution
Q = np.array([0.2, 0.25, 0.3, 0.15, 0.1])

By plotting and the mismatch is immediately visible.

Show code 
import matplotlib.pyplot as plt

x = np.arange(len(vocab))
width = 0.35

fig, ax = plt.subplots(figsize=(7, 4))
ax.bar(x - width/2, P, width, label="P (true)", color="steelblue")
ax.bar(x + width/2, Q, width, label="Q (model)", color="tomato")

ax.set_xticks(x)
ax.set_xticklabels(vocab)
ax.set_ylabel("Probability")
ax.set_title(f"P vs Q")
ax.legend()
plt.tight_layout()
plt.show()

Directly applying the discrete formula :

def kl_divergence(p, q):
    return np.sum(p * np.log(p / q))

kl_pq = kl_divergence(P, Q)
print(f"D_KL(P || Q) = {kl_pq:.4f} nats")
D_KL(P || Q) = 0.1874 nats

The KL-divergence is non-zero so there is a difference between the true distribution and the current distribution of our model. The difference between the true distribution and the predicted distribution is obvious.

The order of the distributions is not trivial when computing KL-divergence. This can be shown shown numerically:

kl_qp = kl_divergence(Q, P)

print(f"D_KL(P || Q) = {kl_pq:.4f} nats")
print(f"D_KL(Q || P) = {kl_qp:.4f} nats")
print(f"Symmetric?   {np.isclose(kl_pq, kl_qp)}")
D_KL(P || Q) = 0.1874 nats
D_KL(Q || P) = 0.2147 nats
Symmetric?   False

Summary

KL-divergence can be used answering the question: how different are two probability distributions? The distribution is used, for example, LLM alignment to tell how far away the models predicted vocabulary distribution is from the training data distribution. KL-divergence has a mathematical foundation in information theory.

References

[1]
D. M. Ziegler et al., “Fine-tuning language models from human preferences,” arXiv preprint arXiv:1909.08593, 2019, Available: https://arxiv.org/abs/1909.08593
[2]
L. Ouyang et al., “Training language models to follow instructions with human feedback,” arXiv preprint arXiv:2203.02155, 2022, Available: https://arxiv.org/abs/2203.02155
[3]
I. Goodfellow, Y. Bengio, and A. Courville, Deep learning. MIT Press, 2016. Available: https://www.deeplearningbook.org

Reuse

Citation

BibTeX citation:
@online{bogossian2026,
  author = {Bogossian, Andreas},
  title = {KL {Divergence}},
  date = {2026-03-17},
  url = {https://andreasbogossian.com/posts/KL-divergence/},
  langid = {en}
}
For attribution, please cite this work as:
A. Bogossian, “KL Divergence.” [Online]. Available: https://andreasbogossian.com/posts/KL-divergence/
]]> information theory probabilistic ML llms https://andreasbogossian.com/posts/KL-divergence/ Mon, 16 Mar 2026 22:00:00 GMT