DEV Community: Umesh Tharuka Malaviarachchi

Double Descent: Why Bigger Models Can Generalise Better After Getting Worse

Umesh Tharuka Malaviarachchi — Wed, 28 Oct 2026 03:20:26 +0000

Classical learning theory tells a tidy story. As a model gets more complex, training error falls, test error first falls and then rises, and the best model sits at the bottom of a U-shaped curve. That story is not wrong. It is incomplete.

In modern machine learning, test error often falls, rises sharply to a peak, and then falls again as capacity keeps growing. This second descent is called double descent. It helps explain why very large models, which can fit their training data exactly, sometimes generalise surprisingly well, while a model just large enough to interpolate the data generalises terribly.

This article explains what double descent is, where the peak comes from, what mathematics describes it, and when you should - and should not - expect to see it.

The classical picture: bias, variance, and the U curve

Start with the familiar decomposition. For a prediction problem with input x, true function f, noise, and a learned predictor, expected test error can be split into three parts: squared bias, variance, and irreducible noise.

Bias measures how far the average model, over different training sets, is from the truth. Variance measures how much the fitted model moves when the training set changes. Noise is what no model can remove.

In the classical regime, increasing complexity reduces bias but increases variance. At first the bias reduction wins, so test error falls. Later variance dominates, so test error rises. The sweet spot is the trade-off point.

That picture assumes we stay in the underparameterised regime, where the number of parameters p is smaller than the number of training examples n. The model cannot fit every training point exactly. Once p passes n, the picture changes.

Crossing the interpolation threshold

Define the interpolation threshold as the point where p is roughly equal to n. At this point the model has just enough capacity to fit the training data almost exactly, including its noise.

Near this threshold, the fitted model is extremely sensitive to the training data. A small change in one example can swing the fit a lot, because the model has almost no freedom left: it must thread through nearly every point. Variance explodes, and test error peaks.

Past the threshold, in the overparameterised regime where p is much larger than n, there are many parameter settings that all fit the training data exactly. The learning algorithm does not pick one at random. It picks a particular one, usually the one closest to its initialisation or the one with smallest norm under the optimisation dynamics. That implicit choice acts as a form of regularisation, and it can reduce variance again. Test error descends a second time.

So the full curve is: descent, ascent to a peak at interpolation, then a second descent. That is double descent.

The maths: least squares and the minimum-norm interpolator

The cleanest setting for double descent is linear regression with more features than examples.

Suppose we have n training examples and p features, with p greater than n. Let X be the n by p design matrix and y the vector of targets. There are infinitely many weight vectors w that satisfy Xw = y exactly. Gradient descent starting from zero, run to convergence on the squared loss, converges to the minimum-norm interpolator:

w = X^{T} (X X^{T})^{- 1} y

This is the interpolating solution with the smallest Euclidean norm. The norm matters because, for linear models, a smaller norm typically means a smoother, less wild function, and therefore often better generalisation.

For ridge regression, the solution is:

w_{λ} = (X^{T} X + λ I)^{- 1} X^{T} y

where $λ$ is the regularisation strength and I is the identity matrix. As $λ$ goes to zero in the overparameterised regime, ridge regression tends to the minimum-norm interpolator above. Positive $λ$ shrinks the solution further and can soften or remove the peak.

A useful way to think about the peak is through the singular values of X. Near p = n, the smallest singular values of the data matrix can become very small. Inverting them, as the minimum-norm formula does, amplifies noise enormously. Far past the threshold, the spectrum spreads out, the effective amplification falls, and the interpolator becomes more stable. Random matrix theory makes this precise for idealised Gaussian features: test risk can be written as a function of the ratio p/n, and it diverges at p/n = 1 in the unregularised limit, then decreases as p/n grows.

The same intuition extends, with more work, to random feature models and to wide neural networks in the kernel regime, where the model behaves like kernel regression with the neural tangent kernel. The details differ, but the mechanism is similar: a peak at interpolation, followed by a second descent driven by the implicit regularisation of the fitting procedure.

A practical example: watching the peak appear

You can see double descent without a large neural network. Take a synthetic regression problem, expand the input with random features, and increase the number of features past the number of training examples. Use ridge regression with a tiny regularisation strength.

import numpy as np
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_squared_error

rng = np.random.default_rng(0)
n_train, n_test = 200, 1000
X_train_raw = rng.normal(size=(n_train, 20))
X_test_raw = rng.normal(size=(n_test, 20))
true_w = rng.normal(size=20)
y_train = X_train_raw @ true_w + 0.5 * rng.normal(size=n_train)
y_test = X_test_raw @ true_w + 0.5 * rng.normal(size=n_test)

def random_features(X, p):
    # Random ReLU features: a simple, explicit expansion.
    W = rng.normal(size=(X.shape[1], p)) / np.sqrt(X.shape[1])
    return np.maximum(X @ W, 0)

for p in [10, 50, 100, 190, 200, 210, 300, 500, 1000, 2000]:
    Phi_train = random_features(X_train_raw, p)
    Phi_test = random_features(X_test_raw, p)
    model = Ridge(alpha=1e-6)
    model.fit(Phi_train, y_train)
    err = mean_squared_error(y_test, model.predict(Phi_test))
    print(f"features={p:5d}  ratio p/n={p / n_train:.2f}  test MSE={err:.3f}")

In a typical run, test error falls as p grows toward n, spikes near p = n, and then falls again as p grows well past n. The exact heights depend on the noise level, the feature distribution, and the regularisation strength, which is exactly the point: double descent is a phenomenon of the whole setup, not a property of model size alone.

Try two variations. First, increase the label noise: the peak gets taller. Second, increase the ridge penalty: the peak shrinks and may disappear. Both effects follow from the maths above. Noise is what gets amplified near interpolation, and explicit regularisation competes with the implicit regularisation of the minimum-norm solution.

Why the second descent does not mean "bigger is always better"

Double descent is often misread as a licence to ignore model size. That reading is wrong for several reasons.

First, the second descent is not guaranteed. It depends on the data distribution, the feature map, the noise level, the loss, and the optimiser. With heavy label noise, poor features, or a badly chosen architecture, test error past the threshold can stay high or even rise again.

Second, the peak is real and dangerous. A model near the interpolation threshold can be the worst model in the whole family, even though its training error is near zero. If you tune model size by training error alone, you will walk straight into the peak.

Third, test error is not the only cost. Larger models cost more to train, serve, and maintain, and they can be harder to interpret and debug. A second descent in test error of a few percent may not justify a tenfold increase in compute.

Fourth, double descent in test error does not imply good behaviour elsewhere. Calibration, robustness to distribution shift, fairness across subgroups, and performance on rare cases can all behave differently from average test error. A model past the peak can still fail badly on the examples that matter most.

Finally, modern practice rarely varies only p with everything else fixed. Architecture, data scale, training time, and regularisation all change together. Real scaling curves blend many effects, and attributing all of an improvement to double descent would be a mistake.

Common pitfalls when testing for double descent

Tuning on test data to find the peak. Use a validation set, and fix the evaluation protocol before sweeping model size.
Forgetting regularisation. Even a tiny ridge penalty, weight decay, or early stopping changes the shape of the curve. Report it.
Confusing interpolation with memorisation in a harmful sense. Interpolation means fitting training points exactly. Whether that harms generalisation depends on how the interpolator is chosen.
Sweeping too coarsely. The peak is narrow. If your grid jumps from p well below n to p well above n, you will miss it and conclude the curve is monotone.
Ignoring the optimiser. Different optimisers, initialisations, and learning rates pick different interpolators. The minimum-norm story is specific to gradient descent from zero on squared loss; other settings need their own analysis.

What to take away

The classical U curve describes the underparameterised regime well. Double descent extends it: test error can peak at the interpolation threshold, where the model just barely fits the data, and then fall again as overparameterisation lets the optimiser's implicit bias select a smoother interpolator.

For practice, the lesson is concrete. Do not judge a model by training error, do not assume the worst model is the largest one, and do not assume the largest one is the best. Sweep capacity carefully, keep regularisation explicit, watch the region around p = n, and validate on data the model has never seen.

Bigger is not always better, and smaller is not always safer. The shape of the curve depends on the problem, and the only way to know it is to measure it.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial

Sharpness-Aware Minimization: Why Flat Minima Generalise Better Than Sharp Ones

Umesh Tharuka Malaviarachchi — Wed, 28 Oct 2026 00:12:04 +0000

Sharpness-Aware Minimization: Why Flat Minima Generalise Better Than Sharp Ones

Two neural networks can reach almost exactly the same training loss and yet behave very differently on new data. One generalises well. The other collapses on the first distribution shift, the first noisy label, or the first slightly different test set. A large part of the difference is not the loss value they reached. It is the shape of the loss landscape around that value.

Sharpness-Aware Minimization, usually shortened to SAM, makes that shape an explicit part of training. Instead of asking only "how low is the loss here?", it asks "how low does the loss stay nearby?" A flat minimum, where the loss stays low in a neighbourhood, is preferred over a sharp minimum, where the loss rises steeply as soon as the weights move a little. That single change of objective explains a surprising amount of SAM's practical behaviour, and also its limits.

The intuition: two minima, same loss, different futures

Picture a one-dimensional loss curve with two valleys. Both reach the same depth. One valley is narrow and steep. The other is wide and gentle. On the training data, the two solutions look identical: same loss, same accuracy.

Now imagine that the test data are not exactly the training data. In effect, the loss landscape shifts slightly, or the weights are perturbed slightly, because the model must work on inputs it has never seen. In the sharp valley, a small shift pushes the model up a steep wall, and the loss explodes. In the flat valley, the same shift moves the model along a gentle floor, and the loss barely changes.

This is the core intuition behind flat minima. Generalisation is not only about fitting the training set. It is about robustness to the small differences between training and test conditions. A solution that is stable under perturbation is, in this sense, a safer bet than a solution that is only good at one exact point.

That intuition needs care. Flatness alone is not a complete theory of generalisation. As discussed below, flatness can be redefined by rescaling parameters, and a flat minimum in one parameterisation can look sharp in another. SAM is best understood as a useful training bias, not as a law of nature.

From intuition to an objective

Let $w$ be the model parameters and $L (w)$ be the training loss. Standard training minimises $L (w)$ directly. SAM instead minimises the worst-case loss in a neighbourhood of $w$ :

L_{SAM} (w) = ∣ ϵ ∣ \leq ρ max L (w + ϵ)

Here $ϵ$ is a perturbation of the weights, $∣ ϵ ∣$ is its size, and $ρ$ is the neighbourhood radius. The parameter $ρ$ controls how far SAM looks. If $ρ$ is tiny, SAM is almost the same as standard training. If $ρ$ is huge, SAM may sacrifice training fit to chase an extremely flat region that is not actually useful.

The minimisation problem therefore becomes:

w min ∣ ϵ ∣ \leq ρ max L (w + ϵ)

This is a minimax problem. The inner maximisation finds the worst nearby perturbation. The outer minimisation moves the weights so that even that worst nearby point has low loss.

Solving the inner problem exactly would be expensive. SAM uses a first-order approximation that is cheap and works well in practice. Expand $L (w + ϵ)$ to first order:

L (w + ϵ) \approx L (w) + ϵ^{⊤} \nabla L (w)

Maximising this linear approximation over $∣ ϵ ∣ \leq ρ$ gives a closed form. With the Euclidean norm, the worst perturbation points in the direction of the gradient:

ϵ^= ρ \frac{\nabla L ( w )}{∣\nabla L ( w ) ∣}

The SAM gradient is then the gradient evaluated not at $w$ , but at the perturbed point $w + ϵ^$ :

\nabla L_{SAM} (w) \approx \nabla L (w + ϵ^)

In words: take a step in the gradient direction to find the worst nearby point, compute the gradient there, and use that gradient to update the original weights. A sharp minimum will have a large gradient at the perturbed point, because the loss rises steeply nearby, so the update will push away from it. A flat minimum will have a small gradient at the perturbed point, so the update will be gentle and the solution will stay put.

The algorithm, step by step

One SAM update, in its simplest form, looks like this:

Compute the gradient $\nabla L (w)$ at the current weights, on the current mini-batch.
Form the perturbation $ϵ^= ρ \nabla L (w) /∣\nabla L (w) ∣$ .
Compute the gradient at the perturbed weights, $\nabla L (w + ϵ^)$ , on the same mini-batch.
Update the weights using that second gradient, with your usual optimiser step, for example stochastic gradient descent with momentum or Adam.

That is the whole idea. The cost is clear: SAM needs two gradient computations per update instead of one, so each step costs roughly twice as much as a standard step. In practice, SAM often needs fewer steps to reach a given test accuracy, or reaches a better final accuracy for the same number of epochs, but the per-step cost is real and must be counted honestly when comparing methods.

A minimal PyTorch-style sketch makes the two-pass structure explicit:

# w are the model parameters, rho is the neighbourhood radius.
# loss_fn returns the training loss on the current mini-batch.

def sam_step(model, optimiser, loss_fn, rho):
    # First pass: gradient at the current weights.
    loss = loss_fn(model)
    loss.backward()
    grads = [p.grad for p in model.parameters()]

    # Form the worst-case perturbation in the gradient direction.
    norm = sum((g ** 2).sum() for g in grads) ** 0.5
    with torch.no_grad():
        for p, g in zip(model.parameters(), grads):
            p.add_(rho * g / (norm + 1e-12))

    # Second pass: gradient at the perturbed weights.
    optimiser.zero_grad()
    loss_perturbed = loss_fn(model)
    loss_perturbed.backward()

    # Restore the original weights, then step with the perturbed gradient.
    with torch.no_grad():
        for p, g in zip(model.parameters(), grads):
            p.sub_(rho * g / (norm + 1e-12))
    optimiser.step()
    optimiser.zero_grad()

This sketch omits optimiser details, adaptive variants, and distributed-training bookkeeping, but it shows the essential point: SAM is not a new optimiser in the usual sense. It is a wrapper around an existing optimiser that changes which gradient is used for the step.

Why it helps, when it helps

SAM helps most when the gap between training and test performance is driven by sensitivity to perturbation. That includes several familiar settings.

First, noisy labels. A sharp minimum often fits label noise by contorting the decision boundary in a fragile way. A flat minimum is less able to do that, because fitting a noisy point exactly usually requires a sharp feature. By penalising sharpness, SAM makes the model less willing to memorise noise at the expense of the clean signal.

Second, small data and large models. When a model has far more capacity than the data need, there are many solutions with near-zero training loss. Standard training may land in any of them, including a sharp one. SAM biases the search toward the flatter members of that set, which tend to be more stable.

Third, distribution shift. If test inputs differ slightly from training inputs, the effective loss landscape shifts. A flat solution tolerates that shift better, for exactly the valley picture described above. SAM does not solve distribution shift, and it does not replace data augmentation or domain adaptation, but it reduces one source of fragility.

It is worth noting what SAM does not do. It does not change the model architecture, the data, or the loss function itself. It changes the training dynamics. Two runs with the same data and architecture can therefore end in different places purely because one used SAM and the other did not, even though both minimised the same nominal loss.

Limitations and common pitfalls

SAM is not a free improvement, and its failures are instructive.

The radius $ρ$ matters, and there is no universal value. Too small, and SAM does almost nothing. Too large, and it underfits, because it refuses to enter any region that is not extremely flat, even if that region fits the data well. In practice, $ρ$ is tuned on validation data, like a learning rate or weight decay, and the best value depends on the model, the batch size, and the task. Reporting a SAM result without reporting $ρ$ makes the result hard to reproduce.

Sharpness is not invariant to reparameterisation. Rescaling the parameters of a network, for example by scaling one layer up and another down in a way that leaves the function unchanged, can change the measured sharpness without changing the model at all. This means that "flat" is partly a property of the parameterisation, not only of the function. Later variants of SAM address this by using adaptive, scale-aware neighbourhoods, where the perturbation size is scaled per parameter. If you use SAM on a model with very different parameter scales across layers, an adaptive variant is usually the safer choice.

SAM costs twice the gradient computation per step. If you compare SAM for 100 epochs against standard training for 100 epochs, SAM has used roughly twice the compute. A fair comparison either matches total compute, or states the compute difference plainly. Many apparent SAM gains shrink when compute is matched, even if a real gain remains.

SAM interacts with other regularisers. Weight decay, dropout, label smoothing, data augmentation, and SAM all push the model toward simpler or more robust solutions, but they do so in different ways, and stacking all of them at full strength can underfit. When adding SAM to an already heavily regularised pipeline, it is usually better to retune the other regularisers downward rather than simply adding SAM on top.

Finally, SAM does not fix a bad loss, bad data, or a mismatched architecture. If the training labels are systematically wrong, if a key subgroup is missing from the data, or if the model cannot represent the true function, seeking a flat minimum will find a flat version of a bad solution. Flatness is a second-order concern. It improves a solution that is already roughly right. It does not create one from nothing.

Practical takeaways

Use SAM when you have a generalisation gap that looks like fragility: training accuracy is high, test accuracy lags, and small changes in data or seed move test accuracy a lot. That is the signature of a sharp-minimum problem, and SAM directly targets it.

Tune $ρ$ on validation data, and try an adaptive variant if your model has layers with very different scales. Start from the standard SAM radius range used for your architecture family, then move up if the model still overfits sharply, and down if training accuracy drops noticeably.

Count the compute. Budget SAM as two gradient passes per step, and compare it against standard training at matched compute before declaring it a win for your task. If SAM still wins at matched compute, you have a real result. If it only wins because it used twice the compute, say so.

Do not treat flatness as a certificate. A flat minimum is a good sign, not a guarantee. Validate on held-out data, check calibration and subgroup performance, and test robustness to the shifts you actually care about. SAM improves the odds. It does not remove the need to measure.

The takeaway

Standard training asks for a point with low loss. SAM asks for a neighbourhood with low loss. That change, implemented as a gradient step at a perturbed point, biases training toward flat minima that tolerate the small differences between training and test conditions. It costs an extra gradient pass, it needs its radius tuned, and its notion of sharpness depends on parameterisation, so it is not a cure-all. But when generalisation fails because the solution is fragile rather than simply wrong, seeking flatness is one of the most direct and interpretable fixes available. The loss value tells you where you are. The shape around it tells you how long you are likely to stay there.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial

Grokking: Why Neural Networks Suddenly Generalise Long After Memorising

Umesh Tharuka Malaviarachchi — Tue, 27 Oct 2026 20:13:23 +0000

Grokking: Why Neural Networks Suddenly Generalise Long After Memorising

Train a small neural network on modular arithmetic and watch the curves. Training accuracy hits 100 percent quickly. Test accuracy sits near chance for a long time. Then, after thousands more steps with no change in the training data, test accuracy jumps to near 100 percent as well. Nothing in the data changed. Nothing in the task changed. The network grokked.

Grokking is the delayed, sudden transition from memorisation to generalisation. It was first noticed clearly on small algorithmic tasks, but the mechanism behind it - a slow competition between a memorising solution and a generalising solution, decided by regularisation and optimisation dynamics - matters far beyond toy arithmetic. It is a warning against judging a model too early, and a window into how networks build internal structure.

The setup: a task that can be memorised or understood

The classic grokking setup is modular addition. Inputs are pairs (a, b), the target is (a + b) mod p, for some prime p. The training set contains only a fraction of all possible pairs, and the test set contains the rest. A model that memorises the training pairs will get training accuracy right and test accuracy wrong. A model that learns the actual addition rule will get both right.

The point of choosing such a clean task is that generalisation is unambiguous. There is a correct rule, it is simple, and it is either learned or it is not. In messier real tasks, memorisation and generalisation blend, which makes the transition harder to see but no less real.

A typical training objective adds weight decay to the usual loss:

L (θ) = L_{data} (θ) + λ ∣ θ ∣^{2}

Here theta is the vector of parameters, L_data is the training loss, and lambda controls the strength of weight decay. That small second term turns out to be central to grokking, as we will see.

What the curves actually look like

In the early phase, the network memorises. It fits the training pairs by building a lookup-like mapping. Training loss falls to near zero. Test loss stays high. To a practitioner watching only the first part of training, the model looks like a failure: it has overfitted, and further training looks pointless.

In the middle phase, almost nothing visible happens on the test curve. Training accuracy stays at 100 percent. Test accuracy stays low. But inside the network, something is changing slowly. The weights are drifting, under the pressure of weight decay, toward a simpler configuration that also fits the training data.

In the late phase, test accuracy rises sharply, often over a relatively small number of steps compared with the long plateau before it. Training accuracy does not change, because it was already perfect. The model has not learned new data. It has reorganised what it already knew.

That shape - long plateau, then sharp jump - is the signature of grokking. It is not gradual improvement. It is a phase transition in the learned solution.

Two solutions, one loss landscape

Think of the network as having (at least) two kinds of solution available.

The memorising solution fits each training pair individually. It needs many effective parameters, because each pair is handled more or less separately. Its weight norm tends to be large: memorisation is expensive in norm.

The generalising solution implements the underlying rule. For modular addition, that rule can be represented with trigonometric features: the network can learn to map a and b to rotations, add the angles, and read off the result. That solution is compact. Its weight norm tends to be smaller, because one shared mechanism replaces many separate memories.

Early in training, gradient descent finds the memorising solution first because it is easier to reach. Fitting training points directly is a straightforward direction for the gradient. Building a clean trigonometric circuit takes longer: it requires coordinating many weights into a coherent structure, and the gradient signal for that structure is weaker while memorisation already gives near-zero loss.

But once training loss is near zero, the data term of the loss stops pushing hard. Weight decay keeps pushing, slowly, toward smaller norm. The memorising solution, with its large norm, is penalised more than the generalising solution. Over many steps, the network drifts from the expensive solution to the cheap one, while staying on the manifold of near-zero training loss. When the drift crosses a threshold, the generalising circuit becomes strong enough to dominate predictions on unseen pairs, and test accuracy jumps.

In short: optimisation finds the fast solution first, regularisation selects the simple solution eventually, and grokking is the time between those two events.

Why weight decay matters so much

Remove weight decay, or make it too small, and grokking may never happen, or may take far longer. The memorising solution has no reason to move. With strong weight decay, the pressure toward low norm is stronger, and grokking tends to happen earlier, though too much decay can also prevent fitting in the first place.

This also explains why grokking time is so sensitive to hyperparameters. Learning rate, initialisation scale, weight decay, and the fraction of data used for training all shift the balance between the two solutions and the speed of the drift. A small change can turn a run that groks in a few thousand steps into one that never groks within the training budget. That sensitivity is not noise to be averaged away. It is evidence that grokking is a competition with a sharp threshold, not a smooth learning curve.

A related quantity is the effective complexity of the solution. One way to formalise the intuition is to compare solutions by a norm or by a description length: the generalising solution is simpler in the sense that it can be described by a shared rule rather than a list of cases. Weight decay is a crude proxy for that simplicity, but in this setting it is a surprisingly effective one.

What is happening inside the network?

For modular addition, researchers have reverse-engineered the generalising circuit. The network learns embeddings that behave like Fourier features: each input is mapped to a set of sine and cosine components, the components are combined in a way that corresponds to adding angles, and a readout picks out the correct residue.

The memorising phase does not build that circuit. It builds something closer to a table. The transition to generalisation corresponds to the Fourier-like structure emerging and taking over, while the table-like components shrink under weight decay.

You do not need to follow every detail of that circuit to take the practical lesson. The lesson is that test accuracy can be flat while internal structure is changing in a way that matters enormously. A flat test curve is not proof that nothing is being learned. It may mean that the learned structure is not yet strong enough to change predictions, even though it is growing steadily underneath.

This is a useful corrective to the habit of early stopping at the first sign of overfitting. Early stopping is often right, but grokking shows a case where it would be exactly wrong: stopping at the plateau would freeze the memorising solution in place and discard a generalising solution that was on its way.

A minimal example you can run

The following PyTorch sketch trains a small model on modular addition. It is deliberately small: a two-layer network on one-hot inputs is enough to show the pattern on a reduced problem, and a small transformer shows it more cleanly on the full task. Treat this as a starting point for experiment, not as a tuned recipe, because grokking is hyperparameter sensitive.

import torch
import torch.nn as nn

p = 97
torch.manual_seed(0)

# All pairs (a, b), target (a + b) mod p
a = torch.arange(p).repeat_interleave(p)
b = torch.arange(p).repeat(p)
y = (a + b) % p
perm = torch.randperm(p * p)
train_idx = perm[:int(0.5 * p * p)]

class SmallMLP(nn.Module):
    def __init__(self, p, hidden=128):
        super().__init__()
        self.emb_a = nn.Embedding(p, hidden)
        self.emb_b = nn.Embedding(p, hidden)
        self.fc = nn.Sequential(
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, p),
        )
    def forward(self, a, b):
        h = self.emb_a(a) + self.emb_b(b)
        return self.fc(h)

model = SmallMLP(p)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1.0)
loss_fn = nn.CrossEntropyLoss()

for step in range(20000):
    idx = train_idx[torch.randint(len(train_idx), (512,))]
    logits = model(a[idx], b[idx])
    loss = loss_fn(logits, y[idx])
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 2000 == 0:
        with torch.no_grad():
            train_acc = (model(a[train_idx], b[train_idx]).argmax(-1) == y[train_idx]).float().mean()
            test_mask = torch.ones(p * p, dtype=torch.bool)
            test_mask[train_idx] = False
            test_acc = (model(a[test_mask], b[test_mask]).argmax(-1) == y[test_mask]).float().mean()
        print(f"step={step} train_acc={train_acc:.3f} test_acc={test_acc:.3f}")

If you run a version of this, do not expect grokking on a fixed schedule. Vary weight decay and the training fraction, and log test accuracy over a long horizon. The interesting runs are the ones where test accuracy stays low for a long stretch and then jumps. If test accuracy rises immediately, you are likely in a regime where the generalising solution is found first, which is good for practice but does not show grokking.

Does grokking happen on real tasks?

This is the right question, and the honest answer is: not usually in this clean form, but the underlying mechanism shows up in related ways.

Real datasets are noisy, tasks are less cleanly separable into memorise versus understand, and training budgets, regularisation, and data scale all differ. You should not expect a single sharp jump in test accuracy on a large language model in the same way you see it on modular addition.

What does carry over is the idea of delayed generalisation through simplification. Several related phenomena point the same way:

Longer training with strong regularisation can improve generalisation after training accuracy has saturated, because the solution keeps simplifying.
Representations can become cleaner and more structured late in training, even when accuracy metrics look flat.
Models can switch from a shortcut solution to a more robust solution when the shortcut is penalised, for example by weight decay, data augmentation, or a change in the training distribution.

In other words, grokking is an extreme, legible case of a broader pattern: the solution found first is not necessarily the solution that survives regularisation, and the difference between them may only appear after a long delay.

Limitations and common pitfalls

Do not generalise from one toy task. Modular addition is chosen because it is clean. Most real tasks do not have a single compact rule waiting to be found, and memorisation may be partly necessary, for example for facts that cannot be derived from a rule.

Do not assume longer training always helps. Grokking is a case where it does, but in many settings longer training after saturation just overfits further or wastes compute. The difference is whether a simpler, equally fitting solution exists and whether your regularisation pushes toward it.

Weight decay is not magic. It helps here because the generalising solution has smaller norm. In other settings, the robust solution may not be the smallest-norm one, and weight decay may hurt. Test the mechanism, do not just copy the hyperparameter.

A flat test curve is ambiguous. It can mean the model is stuck memorising, or it can mean a generalising structure is growing underneath. You cannot tell from the curve alone. Inspect representations, probe the model on structured test cases, or vary regularisation to see which explanation fits.

Grokking time is fragile. Reporting that a model groks without reporting weight decay, initialisation, training fraction, and optimiser settings is incomplete. Small changes move the transition by orders of magnitude, so reproducibility needs the full setup.

Practical takeaways

If training accuracy saturates and test accuracy is flat, do not immediately conclude the run is done. Check whether a simpler solution could exist, and whether regularisation is still pushing the weights.
Log internal structure, not just accuracy, on long runs: weight norms, representation geometry, and performance on structured probes can move before test accuracy does.
Treat weight decay and other regularisers as part of the learning dynamics, not just as a way to prevent overfitting. They decide which of several fitting solutions survives.
Be careful with early stopping on algorithmic or highly structured tasks. The best solution may arrive long after the training loss has bottomed out.
When you see a sudden jump in generalisation, ask what changed underneath: a new circuit, a norm collapse, or a shift in representation. The jump is a symptom, and the mechanism is the interesting part.

Grokking is a reminder that learning is not just about reducing loss. It is about which solution, among the many that achieve low loss, the training process ends up with, and how long it takes to get there. On modular arithmetic, that process is visible as a long plateau followed by a sudden jump. On real tasks, it is quieter and messier, but the same question applies: is your model memorising the training set, or has it found a rule that will survive new data? Sometimes you have to wait a surprisingly long time for the answer.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial

Federated Learning: Why Averaging Model Updates Works - and When It Fails

Umesh Tharuka Malaviarachchi — Tue, 27 Oct 2026 05:03:33 +0000

Most machine learning still assumes a convenient fiction: all the data sits in one place, ready to be shuffled, split, and trained on. Real data rarely cooperates. It lives on phones, in hospitals, in banks, and in factories, where moving it is expensive, slow, or simply not allowed. Federated learning is the attempt to train anyway, without centralising the data.

The core idea is disarmingly simple. Instead of sending data to the model, send the model to the data. Each client trains locally, sends back only an update, and a server averages those updates into a new global model. That averaging step, called Federated Averaging or FedAvg, looks almost too naive to work. Understanding why it usually does work, and exactly where it breaks, is the key to using it well.

The problem FedAvg solves

Suppose there are K clients. Client k holds n_k examples, and n is the total number of examples across all clients. The global objective is a weighted average of local objectives:

F (w) = k = 1 \sum K \frac{n _{k}}{n} F_{k} (w)

Here w is the model, and $F_{k} (w)$ is the loss on client k's data alone. The server never sees the data behind $F_{k}$ . It only sees updates derived from it.

In a centralised setting, stochastic gradient descent would repeatedly compute a gradient on a random batch drawn from the whole population. In federated learning, each client can only draw batches from its own slice. If every client held a representative sample of the whole population, this would barely matter. In practice, clients are heterogeneous: one hospital sees different patients from another, one phone user types different words from the next, and one sensor sees a different operating regime. That heterogeneity, usually called non-IID data, is the central difficulty of the field.

FedAvg in one round

A round of FedAvg proceeds as follows:

The server sends the current global model $w_{t}$ to a subset of clients.
Each selected client runs several steps of SGD on its own data, starting from $w_{t}$ , and obtains a local model $w_{t k}$ .
Each client sends its update, or its new model, back to the server.
The server forms the next global model as a weighted average:

w_{t + 1} = k \sum \frac{n _{k}}{n} w_{t k}

Two design choices matter enormously here. First, clients do multiple local steps rather than one. That reduces communication, which is often the real bottleneck when clients are phones on slow networks or devices that are only intermittently available. Second, the averaging is weighted by dataset size, so a client with more data pulls the global model more strongly. Both choices are sensible defaults, and both can mislead.

Why should averaging models work at all? Averaging parameters is not the same as averaging predictions. For a linear model, the two coincide in a useful way. For a deep network, averaging two models trained from different starting points usually produces a poor model, because the same function can be represented by permuted hidden units. FedAvg avoids the worst of this by starting every client from the same $w_{t}$ each round and only letting them drift for a few local steps. The local models stay in the same basin, so their average remains a sensible model. Stretch the local training too far, and that assumption fails.

A minimal FedAvg loop

The following sketch captures the structure, stripped of networking, serialisation, and client selection. It is useful precisely because it shows how little code the core idea needs, and how much real systems add around it.

# w: global model parameters (a flat vector for clarity)
# clients: list of local datasets, each with its own loss and gradient
def client_update(w, data, steps, lr):
    local_w = w.copy()
    for _ in range(steps):
        batch = sample_batch(data)
        local_w = local_w - lr * gradient(local_w, batch)
    return local_w

def federated_round(w, clients, weights, steps, lr):
    updates = []
    for client, weight in zip(clients, weights):
        updates.append(weight * client_update(w, client, steps, lr))
    return sum(updates)  # weights sum to 1, so this is the weighted average

In a real deployment, client_update runs on the device, federated_round runs on the server, and only the returned vector crosses the network. Everything else - dropout handling, secure aggregation, versioning, retries for clients that vanish mid-round - is engineering around this loop.

Client drift: the failure mode to watch

When data are heterogeneous, each client's local optimum is different. Let $ParseError: KaTeX parse error: Expected group after '^' at position 4: w_k^̲$ be the minimiser of $F_{k}$ . Local SGD pulls client k toward $ParseError: KaTeX parse error: Expected group after '^' at position 4: w_k^̲$ , not toward the global minimiser. After several local steps, the average of the local models can drift away from a good global solution, even if each local step reduced its own client's loss.

A simple picture helps. Imagine two clients, one whose data are mostly cats and one whose data are mostly dogs. Each local model becomes excellent at its own class and worse at the other. Their average may be mediocre at both. This is not a bug in the averaging formula. It is what the formula does when the local objectives disagree.

Several symptoms point to client drift in practice:

Training loss falls on each client but global validation accuracy stalls or oscillates.
More local steps make things worse, not better, even though they reduce communication.
A small subset of clients dominates, because they have more data or are selected more often.
Performance varies wildly across clients after training, even when average accuracy looks acceptable.

The standard remedies trade off against each other. Fewer local steps reduce drift but increase communication. A proximal term that penalises movement away from $w_{t}$ , as in FedProx, keeps local models closer to the shared starting point. Momentum-like corrections and control variates, as in SCAFFOLD, explicitly estimate and subtract the drift. None of these removes heterogeneity. They manage it, at the cost of extra computation, extra state, or extra tuning.

Communication, computation, and stragglers

Federated learning inverts the usual resource picture. In a datacentre, compute is the bottleneck and communication is cheap. On the edge, communication is expensive, devices are slow and heterogeneous, and some clients simply disappear. A round waits for the slowest selected client unless the system imposes a deadline, and a deadline means some updates are discarded.

This creates three practical trade-offs:

Local steps vs. rounds. More local work per round saves bandwidth but increases drift and straggler delay.
Clients per round vs. stability. Sampling more clients gives a better estimate of the global gradient but costs more coordination.
Model size vs. participation. A larger model may be more accurate, but every extra parameter must be sent down and back up, which excludes slow or constrained clients.

Compression, quantisation, and partial participation all help, but they change the optimisation dynamics. A compressed update is a noisy update. The question is never whether compression loses information. It does. The question is whether the lost information mattered for the global objective, and that depends on how redundant the clients' updates already are.

Privacy: what averaging does and does not guarantee

It is tempting to say that federated learning is private because raw data never leaves the device. That claim is too strong. Model updates are functions of the data, and in some settings they leak a surprising amount about it. Gradient inversion attacks can reconstruct recognisable inputs from gradients for small batches and certain architectures. Membership inference can reveal whether a particular example was in a client's training set. A curious server, or a compromised aggregator, is a real threat model, not a theoretical one.

This is why production systems layer additional protections on top of averaging:

Secure aggregation lets the server learn only the sum of updates, not any individual client's update, using cryptographic masking that cancels out in the aggregate.
Differential privacy adds calibrated noise to updates, so that the influence of any single example is bounded. The cost is accuracy, and the cost grows when many rounds are run.
Data minimisation on the client, such as training on features rather than raw inputs where possible, reduces what any leaked update can reveal.

None of these is free, and none should be treated as a checkbox. Differential privacy in particular forces a concrete choice: how much accuracy are you willing to trade for a stated privacy guarantee, measured over the full training run, not just one round? If a proposal cannot answer that, its privacy claim is marketing, not engineering.

When federated learning is the right tool

Federated learning earns its complexity when at least one of these holds:

Data cannot legally or contractually leave its location, as with health records or financial transactions.
Data are too large or too dispersed to centralise, as with sensor networks.
The deployment itself is distributed, so a model trained centrally would face a distribution shift on the devices that actually use it.

It is a poor fit when data could be centralised with modest effort, when clients are nearly homogeneous, or when the team cannot afford the systems work: client management, versioning, monitoring per-client performance, and debugging a training run where the data are, by design, invisible to the developers. In those cases, centralised training with strong access controls is simpler, faster, and easier to audit.

A useful test before committing: simulate federation on data you can already centralise, by partitioning it artificially across clients in a way that mimics the real heterogeneity. If FedAvg underperforms centralised training badly in simulation, it will not magically improve when the data become genuinely inaccessible. If it performs close to centralised training, the remaining work is systems work, not algorithmic hope.

Common pitfalls

Treating the average as the whole story. Report per-client performance, not just the global average. A model that works for the majority and fails for a minority is a deployment failure, not a success.
Tuning on IID partitions. Randomly shuffling centralised data across clients hides exactly the heterogeneity that makes federation hard. Partition by natural client boundaries, or by label skew, to get an honest estimate.
Ignoring client availability. Clients that are online at night differ systematically from clients online during the day. Selection bias in federated learning is not a sampling detail. It changes the objective being optimised.
Adding privacy after the fact. Noise and secure aggregation interact with optimisation. They need to be part of the training design and evaluation, not bolted on before launch.
Assuming more rounds fix drift. If local objectives disagree, more rounds of the same procedure can converge to a compromise that satisfies nobody. Diagnose drift first, then choose a correction, rather than simply training longer.

Takeaway

Federated averaging works because it exploits a narrow but powerful regime: clients start from a shared model, move only a little on their own data, and are then pulled back together by averaging. That regime gives communication efficiency without centralising data. It breaks when local steps go too far, when clients disagree too much, or when privacy is assumed rather than engineered.

The art of federated learning is therefore not in the averaging formula itself. It is in choosing how far clients may drift before they are reined in, how much noise privacy requires, and how honestly per-client performance is measured. Get those choices right, and averaging becomes a practical way to learn from data that can never be gathered in one place. Get them wrong, and the average is just a polite way of describing a model that works nowhere in particular.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial

Graph Neural Networks: Why Message Passing Works - and Where It Breaks

Umesh Tharuka Malaviarachchi — Tue, 27 Oct 2026 02:20:00 +0000

Graph Neural Networks: Why Message Passing Works - and Where It Breaks

Most machine learning assumes a table. Rows are independent, columns are fixed, and shuffling the rows changes nothing except the order of the output. Graphs break all three assumptions at once. A molecule is not a row. A social network is not a fixed-size vector. A citation graph has no natural ordering of neighbours, and the number of neighbours changes from node to node.

Graph neural networks (GNNs) are the attempt to build a neural model that respects that structure. Their core idea is disarmingly simple: a node learns about itself by listening to its neighbours, then those neighbours have already listened to theirs. Repeat that a few times, and every node carries a summary of its local neighbourhood. Almost everything interesting about GNNs, including their failures, follows from that one mechanism.

The intuition: representation by neighbourhood

Think of a single paper in a citation network. Its own features might be a bag of words from its abstract. That is useful, but it says nothing about the field the paper belongs to, the methods it builds on, or the community that cites it. Its neighbours say a great deal. If most of its neighbours are about protein folding, the paper is probably about protein folding too, even if its abstract is ambiguous.

Message passing turns that intuition into computation. At each layer, every node collects messages from its neighbours, combines them into one aggregated message, and updates its own representation using that aggregate and its previous representation. After one layer, a node knows about its immediate neighbours. After two layers, it knows about neighbours of neighbours. After k layers, it knows about its k-hop neighbourhood.

Crucially, the order in which neighbours send their messages must not matter. A graph has no canonical ordering of neighbours, so aggregation must be permutation invariant: sum, mean, and max all qualify, while concatenating neighbours in a fixed order does not. This is not a cosmetic constraint. It is what makes the model a function of the graph rather than a function of an arbitrary labelling of the graph.

The formalism: one layer, written out

Let a graph have node features
$x_{v}$
for each node
$v$ ,
and let
$N (v)$
be the set of neighbours of
$v$ .
A general message passing layer updates the representation
$h_{v}$
as follows. First, each neighbour sends a message:

m_{uv} = ϕ (h_{u}, h_{v}, e_{uv})

where
$e_{uv}$
is an optional edge feature and
$ϕ$
is a learnable message function. Then the messages are aggregated:

m_{v} = u \in N (v) ⨁ m_{uv}

where
$⨁$
is a permutation invariant aggregator such as sum or mean. Finally, the node updates its state:

h_{v'} = ψ (h_{v}, m_{v})

where
$ψ$
is a learnable update function, usually a small neural network.

The best known special case is the graph convolutional network (GCN). In matrix form, one GCN layer is:

H^{'} = σ (D ~^{- 1/2} A ~ D ~^{- 1/2} H W)

Here
$H$
holds the node representations as rows,
$A ~ = A + I$
is the adjacency matrix with self-loops added so a node also listens to itself,
$D ~$
is the corresponding degree matrix,
$W$
is a learnable weight matrix, and
$σ$
is a nonlinearity. The symmetric normalisation
$D ~^{- 1/2} A ~ D ~^{- 1/2}$
is doing real work: without it, high-degree nodes would receive messages with much larger magnitudes than low-degree nodes, and training would be dominated by hubs.

That single equation explains both the power and the limitation of the GCN. It is a form of Laplacian smoothing: each layer replaces a node representation by a weighted average of its neighbourhood. Smoothing is exactly what you want when neighbours tend to share labels, and exactly what you do not want when they do not.

A practical example

Consider node classification on a citation graph: predict the topic of each paper. A minimal GCN in PyTorch Geometric style looks like this:

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GCN(torch.nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super().__init__()
        self.conv1 = GCNConv(in_dim, hidden_dim)
        self.conv2 = GCNConv(hidden_dim, out_dim)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

Two layers are not a timid choice here. They are often the best choice, for reasons given below. Train this on a labelled subset of nodes, evaluate on held-out nodes, and on a homophilic citation graph it will usually beat a model that uses node features alone by a clear margin. The gain comes from structure, not from extra parameters: the graph itself is doing part of the work.

A common variant worth knowing is GraphSAGE, which samples a fixed number of neighbours instead of using all of them. That matters in practice because real graphs have hubs with thousands of neighbours, and a full neighbourhood aggregation at every layer quickly becomes a bottleneck. Sampling trades a little information for a model that can actually be trained on large graphs. Graph attention networks (GATs) take a different route: they learn a weight for each neighbour, so the model can decide that some neighbours matter more than others. Attention helps, but it does not remove the structural limits discussed next.

Where message passing breaks

1. Oversmoothing: depth destroys information

Each message passing layer averages a node with its neighbours. Repeat that many times and every node starts to look like every other node. Formally, repeated Laplacian smoothing drives representations toward a stationary distribution that depends mainly on node degree, not on the original features. In practice, GCN accuracy often peaks at two or three layers and then falls as depth increases.

This is counterintuitive if you come from images or text, where deeper is usually better. On graphs, depth has a different meaning: depth is receptive field size, and a large receptive field on a densely connected graph is nearly the whole graph. If every node can see the whole graph, and every node averages what it sees, then every node ends up with a similar summary.

Practical responses include residual connections, jumping knowledge networks that combine representations from all layers instead of only the last one, and normalisation schemes designed for graphs. None of these fully removes the problem. The honest rule is to treat depth as a hyperparameter to tune, not as a virtue, and to be suspicious of any graph model with many layers and no explicit mechanism against smoothing.

2. Oversquashing: long-range information is squeezed through a bottleneck

Oversmoothing is about too much mixing. Oversquashing is about too little bandwidth. Suppose a node needs information from a node ten hops away. That information must travel along paths, and at each step it is compressed into a fixed-size vector together with messages from every other neighbour. By the time it arrives, it has been diluted, overwritten, or both.

The effect is worst on graphs with bottlenecks: two dense communities connected by a single bridge edge. All cross-community information must pass through that bridge, and a fixed-size message cannot carry it all. This is a structural property of the graph, not a training failure, and adding layers makes it worse rather than better, because longer paths mean more compression steps.

Rewiring methods, which add or reweight edges to shorten important paths, and graph transformers, which allow direct attention between distant nodes, are both attempts to relieve this bottleneck. They help, but they change the inductive bias of the model: a transformer that attends to all nodes is no longer strictly a local message passing model, and that trade-off should be made deliberately.

3. Heterophily: when neighbours disagree

Message passing assumes, implicitly, that neighbours are informative about each other. On homophilic graphs, where connected nodes tend to share labels, that assumption holds. On heterophilic graphs, where connected nodes tend to differ, it fails. A classic example is a dating network or a fraud network: fraudsters deliberately connect to legitimate users, so averaging a node with its neighbours mixes the signal you are trying to detect with the noise you are trying to avoid.

The fix is not to abandon graphs, but to stop assuming that aggregation should always smooth. Models for heterophilic graphs separate a node's own features from its neighbours' features, use signed or higher-order aggregations, or learn when to listen and when to ignore. The diagnostic step comes first: measure the label agreement between connected nodes before choosing a model. If neighbours usually disagree, a plain GCN is the wrong baseline.

4. Expressivity: the Weisfeiler-Lehman limit

There is a formal ceiling on what standard message passing can distinguish. Message passing GNNs are at most as powerful as the Weisfeiler-Lehman (WL) graph isomorphism test at telling graphs apart. Two graphs that the WL test cannot distinguish will also fool a standard GNN, no matter how large it is.

Concretely, a GNN that aggregates only neighbour multisets cannot count certain structures, such as triangles in some configurations, and cannot distinguish some regular graphs whose nodes all have identical neighbourhoods. This matters less for node classification on a single large graph, where features break the symmetry, and more for graph classification, where the whole graph is the input and structure is the signal. Higher-order GNNs and subgraph-based methods push past the WL limit, at a computational cost. Know which side of that limit your task sits on before promising that a bigger GNN will solve it.

Common pitfalls in practice

Leaking labels through the graph. In transductive node classification, the test nodes are present in the graph during training, just unlabelled. That is standard, but it means random node splits are not the same as random graph splits. If your deployment will see an entirely new graph, evaluate on a new graph, not on held-out nodes of the same graph.
Ignoring degree. A model that sums neighbour messages without normalisation will be dominated by hubs. A model that averages will wash out the signal from a single important neighbour. Check both, and check what happens to low-degree and high-degree nodes separately, not just the average accuracy.
Treating edges as uninformative. Many real graphs have edge types, weights, or timestamps. A model that ignores them is throwing away structure. If edges carry the signal, use a model that conditions messages on edge features, as in the general form above.
Scaling by stacking layers. Neighbourhoods grow exponentially with depth on many graphs. A three-layer GNN on a graph with average degree ten already touches up to a thousand nodes per output. Sample neighbours, use mini-batch training, or accept a shallow model. Memory blow-ups in GNN training are usually neighbourhood explosion, not model size.
Reporting only average accuracy. GNN failures are rarely uniform. Oversmoothing hurts minority classes first, heterophily hurts specific node types, and bottlenecks hurt specific communities. Break results down by degree, by community, and by class before concluding that the model works.

When to use a GNN, and when not to

Use a GNN when the relationships between items carry information that the items alone do not: molecules, citation and social networks, recommender systems viewed as user-item graphs, traffic networks, and knowledge graphs. Use it when the graph is available at prediction time and when the task genuinely depends on neighbourhood structure.

Do not use a GNN when the data are really tabular and the graph was constructed artificially, for example by connecting nearest neighbours in feature space. In that case the graph adds little beyond what a standard model already sees, and it adds failure modes that a standard model does not have. A GNN is not a generic upgrade to a multilayer perceptron. It is a strong inductive bias, and strong biases help only when they match the problem.

The takeaway

Message passing is a simple idea with sharp edges. A node updates itself from its neighbours, the update must not depend on neighbour order, and repeating that update builds a neighbourhood summary. That mechanism explains why GNNs work so well on homophilic graphs, why depth is a liability rather than a virtue, why long-range dependencies get squashed, and why some graphs are fundamentally beyond the reach of standard message passing.

The practical skill is not stacking more layers. It is diagnosing the graph first: how homophilic is it, where are its bottlenecks, how fast do neighbourhoods grow, and what does the task actually require a node to know? Answer those questions, and the choice of GNN, if a GNN is the right choice at all, becomes much clearer.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial

Hamiltonian Monte Carlo: How Momentum Turns Random Walks into Efficient Sampling

Umesh Tharuka Malaviarachchi — Mon, 26 Oct 2026 19:55:05 +0000

Hamiltonian Monte Carlo: How Momentum Turns Random Walks into Efficient Sampling

Most Bayesian models are easy to write down and hard to sample from. You can usually evaluate an unnormalised posterior density, or its gradient, but you cannot draw independent samples from it directly. Markov chain Monte Carlo (MCMC) solves this by constructing a chain whose long-run distribution is the target. The question is how efficiently that chain explores.

Hamiltonian Monte Carlo (HMC) is one of the most effective answers for continuous, differentiable models. Instead of proposing a small random step and hoping it lands somewhere useful, HMC borrows an idea from physics: give the current point momentum, let it glide along the contours of the target distribution, and only then decide whether to accept the move. The result is a sampler that can travel far in one step while still keeping a high acceptance rate.

The problem with random walks

Consider a standard random-walk Metropolis sampler. From the current state, it proposes a new state by adding Gaussian noise, then accepts or rejects that proposal based on the density ratio. This works, but it has a structural weakness: to keep the acceptance rate reasonable, the step size must stay small, and small steps mean the chain diffuses slowly.

In high dimensions this gets worse. A random walk needs on the order of the square of the distance it must travel to cover that distance, because successive steps partly cancel each other out. If the target is strongly correlated, for example a narrow ridge in parameter space, an isotropic proposal will either step off the ridge and be rejected, or crawl along it very slowly.

HMC avoids this by using gradient information to propose coherent, long-distance moves.

A physical analogy that is more than an analogy

Imagine a particle sliding without friction over a landscape. Valleys correspond to regions of high posterior probability, and hills correspond to regions of low probability. If you give the particle a push, it will roll down into a valley, climb partway up the other side, turn around, and keep moving in a way that respects the shape of the landscape.

HMC makes this literal. Let

$q$

denote the position, which is the parameter vector we want to sample, and let

$p$

denote an auxiliary momentum variable of the same dimension. Define a potential energy from the target density

$π (q)$

U (q) = - lo g π (q)

so that high probability means low potential energy. Define a kinetic energy as

K (p) = \frac{1}{2} p^{T} M^{- 1} p

where

$M$

is a symmetric positive definite mass matrix, often taken as the identity or as a diagonal estimate of the posterior scales. The total energy, or Hamiltonian, is

H (q, p) = U (q) + K (p)

The dynamics are given by Hamilton's equations:

\frac{d q}{d t} = \frac{\partial H}{\partial p} = M^{- 1} p, \frac{d p}{d t} = - \frac{\partial H}{\partial q} = - \nabla U (q) = \nabla lo g π (q)

Two properties matter. First, the dynamics conserve total energy, so a trajectory stays on a level set of

$H$

Second, the dynamics preserve volume in

$(q, p)$

space and are reversible in time. Those properties are what allow HMC to propose distant states and still correct exactly for numerical error with a Metropolis step.

The HMC algorithm

One HMC iteration works as follows:

Sample fresh momentum from its Gaussian distribution, independently of the current position. With the kinetic energy above, this is

p \sim N (0, M)

Simulate Hamiltonian dynamics for a fixed trajectory length, starting from the current

$(q, p)$

pair. This produces a proposed pair

$(q, p)$

Accept the proposal with probability

min (1, exp (H (q, p) - H (q, p)))

and otherwise stay at the current position. Finally, discard the momentum: the marginal distribution of

$q$

under the joint distribution is exactly the target, because the joint density is proportional to

$exp (- H (q, p))$

and integrating out a Gaussian momentum leaves

$π (q)$

If the dynamics were simulated exactly, energy would be conserved and every proposal would be accepted. In practice we must discretise, and the Metropolis correction removes the bias introduced by that discretisation.

Why the leapfrog integrator matters

HMC almost always uses the leapfrog integrator, not a generic numerical solver, because leapfrog preserves the two properties that make the correction work: it is volume-preserving and time-reversible. A single leapfrog step of size

$ϵ$

is:

p_{t + ϵ /2} = p_{t} - \frac{ϵ}{2} \nabla U (q_{t})

q_{t + ϵ} = q_{t} + ϵ M^{- 1} p_{t + ϵ /2}

p_{t + ϵ} = p_{t + ϵ /2} - \frac{ϵ}{2} \nabla U (q_{t + ϵ})

Repeat this for

$L$

steps to get a trajectory of total length

$L ϵ$

Leapfrog is symplectic, meaning it approximately preserves the geometric structure of Hamiltonian dynamics, so energy error stays bounded over long trajectories instead of drifting steadily. A Runge-Kutta method, by contrast, is generally not volume-preserving, and using it would break the simple Metropolis correction.

The step size controls accuracy: too large, and energy error grows and proposals are rejected; too small, and computation is wasted on many tiny steps. The trajectory length controls how far the chain travels: too short, and HMC degenerates back toward a random walk; too long, and the trajectory may loop back near its starting point, wasting gradient evaluations. A practical rule is to randomise the trajectory length slightly on each iteration, which avoids periodic behaviour where a fixed length repeatedly returns to almost the same place.

A concrete example: a correlated Gaussian

The advantage of HMC is easiest to see on a strongly correlated two-dimensional Gaussian. A random-walk proposal aligned with the coordinate axes must take tiny steps to stay on the narrow ridge. HMC, by contrast, follows the gradient, so its trajectory curves along the ridge itself.

A minimal leapfrog implementation in NumPy makes the structure clear:

import numpy as np

def leapfrog(q, p, grad_log_pi, epsilon, n_steps):
    # grad_log_pi returns the gradient of log pi at q
    p = p + 0.5 * epsilon * grad_log_pi(q)
    for i in range(n_steps):
        q = q + epsilon * p
        if i != n_steps - 1:
            p = p + epsilon * grad_log_pi(q)
    p = p + 0.5 * epsilon * grad_log_pi(q)
    # Negate momentum for reversibility; harmless for a symmetric proposal
    return q, -p

def hmc_step(q, log_pi, grad_log_pi, epsilon, n_steps):
    p0 = np.random.randn(*q.shape)
    q_new, p_new = leapfrog(q.copy(), p0, grad_log_pi, epsilon, n_steps)
    log_accept = (
        log_pi(q_new) - 0.5 * np.dot(p_new, p_new)
        - (log_pi(q) - 0.5 * np.dot(p0, p0))
    )
    if np.log(np.random.rand()) < log_accept:
        return q_new, True
    return q, False

Here the mass matrix is the identity, so the kinetic energy is simply half the squared norm of the momentum. In a real model, you would replace the identity with a diagonal mass matrix estimated from warmup samples, which rescales coordinates so that a single step size works reasonably well in every direction.

For a Bayesian logistic regression, the same code pattern applies with the log posterior and its gradient in place of the toy density. The gradient is the expensive part, but it is also the reason each iteration can move much further than a gradient-free proposal.

Tuning, warmup, and the No-U-Turn Sampler

In practice, nobody hand-tunes

$ϵ$

and

$L$

for every model. Modern implementations split sampling into a warmup phase and a sampling phase. During warmup, the step size is adapted to hit a target acceptance rate, often around 0.8, and the mass matrix is estimated from the samples seen so far. Those adapted values are then frozen for the sampling phase, so the chain is a valid Markov chain with a fixed transition kernel.

The No-U-Turn Sampler (NUTS), used by default in Stan and PyMC, goes further: it automatically chooses a trajectory length by simulating both forward and backward in time until the trajectory starts to turn back on itself, the "U-turn" condition. This removes the most fragile tuning choice in HMC, at the cost of a more complex algorithm and a variable number of gradient evaluations per iteration. If you use a probabilistic programming library, you are almost certainly using NUTS rather than hand-tuned HMC, but the underlying dynamics are the same.

Diagnostics: what to watch

HMC gives you diagnostic signals that a random-walk sampler does not. The most important is the divergent transition. A divergence occurs when the leapfrog integrator encounters a region of very high curvature, for example the neck of a funnel-shaped posterior, and the energy error becomes large. Divergences are not just rejected proposals: they indicate that the sampler cannot reliably explore that region, and posterior estimates from a chain with many divergences should not be trusted. The usual remedies are a smaller step size, a higher target acceptance rate during warmup, or a reparameterisation of the model.

Beyond divergences, the usual MCMC checks still apply: split-R-hat near 1 across independent chains, effective sample size large relative to the number of draws, and trace plots without long flat stretches or trends. HMC typically achieves a much higher effective sample size per iteration than random-walk Metropolis on smooth, continuous problems, which is why the extra gradient cost is usually worth paying.

Limitations and common pitfalls

HMC is not a universal sampler. It requires a differentiable log density, so it does not apply directly to discrete parameters. A common workaround is to marginalise discrete parameters out when possible, or to use a different sampler for the discrete part in a Gibbs-style scheme.

Multimodal targets are another weakness. Hamiltonian trajectories conserve energy, so a particle starting in one mode with a typical momentum draw will rarely gain enough energy to cross a low-density valley into another mode. Tempering, or a sampler designed for mode jumping, is a better fit there.

Three pitfalls recur in practice. First, forgetting that the gradient must be correct: a wrong sign, a missing Jacobian adjustment after a transformation, or a gradient computed at the wrong point will silently produce a wrong stationary distribution, because the Metropolis correction only corrects integrator error, not gradient error. Second, adapting the step size or mass matrix throughout sampling rather than only during warmup, which breaks the Markov property. Third, interpreting a high acceptance rate as success on its own: a chain with tiny steps accepts almost everything and explores almost nothing. Acceptance rate, effective sample size, and divergences must be read together.

Takeaway

Random-walk Metropolis explores by diffusion, which is slow in high dimensions and on correlated targets. Hamiltonian Monte Carlo explores by dynamics: momentum carries the chain coherently along the contours of the posterior, the leapfrog integrator preserves the geometry that makes the correction exact, and a Metropolis step removes the remaining numerical bias. That combination is why HMC, and its automatic cousin NUTS, became the default engine for modern Bayesian computation on continuous models. It asks for more per iteration, a gradient and a careful warmup, and in return it gives back far more per iteration in effective exploration.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial

Score Matching: Why Learning the Gradient of a Log-Density Beats Learning the Density Itself

Umesh Tharuka Malaviarachchi — Mon, 26 Oct 2026 13:41:30 +0000

Score Matching: Why Learning the Gradient of a Log-Density Beats Learning the Density Itself

Most generative models try to learn a probability density. That sounds like the obvious goal, but it hides an awkward obstacle: the normalising constant. Score matching sidesteps it entirely by learning a different object, the gradient of the log-density, and that single change unlocks a family of models that includes modern diffusion models.

This article explains what the score is, why matching it is a sensible objective, how the practical estimators work, and where the idea breaks down.

The problem with densities

Suppose we want to model data with an unnormalised density

p ~ (x; θ)

The true density is

p (x; θ) = \frac{p ~ ( x ; θ )}{Z ( θ )}

where the partition function is

Z (θ) = \int p ~ (x; θ) d x

For any model beyond the simplest cases, that integral is intractable. Maximum likelihood therefore requires either a model class with a tractable normaliser, such as a normalising flow, or an expensive approximation, such as Markov chain Monte Carlo inside every training step.

Score matching asks a different question: what if we never need the normaliser at all?

The score

The score of a density is the gradient of its log-density with respect to the data:

s (x; θ) = \nabla_{x} lo g p (x; θ)

Here is the key observation. Take the log of the normalised density:

lo g p (x; θ) = lo g p ~ (x; θ) - lo g Z (θ)

The term $lo g Z (θ)$ does not depend on $x$ , so it vanishes under differentiation:

\nabla_{x} lo g p (x; θ) = \nabla_{x} lo g p ~ (x; θ)

The score of the normalised model equals the score of the unnormalised model. We can compute it without ever evaluating the partition function.

Intuitively, the score is a vector field that points uphill, towards regions of higher density. At a mode of the distribution the score is zero. In a valley between two modes it points towards the nearer mode. Learning the score is therefore like learning the shape of the landscape without learning its absolute height.

The Hyvarinen objective

Hyvarinen's score matching objective measures the squared distance between the model score and the true data score, averaged over the data distribution $p_{d a t a}$ :

J (θ) = \frac{1}{2} E x \sim p d a t a [∣ s (x; θ) - \nabla_{x} lo g p_{d a t a} (x) ∣^{2}]

This looks unusable, because it contains the unknown true score $\nabla_{x} lo g p_{d a t a} (x)$ . The breakthrough is an integration-by-parts identity. Under mild regularity conditions, the objective can be rewritten in a form that depends only on the model:

J (θ) = E x \sim p d a t a [tr (\nabla_{x} s (x; θ)) + \frac{1}{2} ∣ s (x; θ) ∣^{2}] + C

where $C$ is a constant independent of $θ$ , and $tr (\nabla_{x} s (x; θ))$ is the trace of the Jacobian of the score, which is the sum of its second derivatives, that is, the Laplacian of the log-density.

So we can train by minimising an expectation over real data, using only quantities we can compute from the model. No samples from the model are needed, and no normaliser is ever evaluated.

Why the trace term matters

The two terms in the objective pull in opposite directions, and both are needed.

The term $\frac{1}{2} ∣ s (x; θ) ∣^{2}$ alone would be minimised by a flat score, a model that says every point is equally likely. The trace term penalises that: where the score changes rapidly, the trace term records curvature. Together they force the model score to match the data score, not merely to shrink to zero.

This is also the main practical weakness of the original objective. The trace of a Jacobian in high dimensions is expensive. For a neural network score model in, say, image space, computing the full Jacobian trace per example is prohibitive. That cost motivated the estimators that followed.

Denoising score matching

Denoising score matching, introduced by Vincent, replaces the intractable trace with a much simpler regression problem. Perturb each data point with noise, for example Gaussian noise:

x ~ = x + σ ϵ, ϵ \sim N (0, I)

The score of the perturbed distribution is known in closed form conditional on the clean point:

\nabla_{x ~} lo g q (x ~ ∣ x) = - \frac{x ~ - x}{σ ^{2}}

So we train the score network to predict the noise direction that would denoise the sample:

J_{D SM} (θ) = E_{x, x ~} [∣ ∣ s (x ~; θ) + \frac{x ~ - x}{σ ^{2}} ∣ ∣^{2}]

This is just supervised regression. No Jacobian trace, no second derivatives, and it scales to high dimensions. When the noise level is small, the score of the perturbed distribution approximates the score of the clean distribution, and the estimator is consistent.

A minimal PyTorch sketch shows the shape of the training loop:

# s_theta: score network, sigma: noise level
for x in dataloader:
    noise = torch.randn_like(x)
    x_noisy = x + sigma * noise
    target = -noise / sigma          # score of q(x_noisy | x)
    pred = s_theta(x_noisy)
    loss = ((pred - target) ** 2).mean()
    loss.backward()
    optimizer.step()

Note what the network actually learns here. It does not learn to output clean data. It learns a vector field: at each noisy point, which direction points back towards the data manifold. That field is exactly what a sampler needs.

Sliced score matching

Sliced score matching offers another way around the trace term. Instead of matching the full score vector, it matches random one-dimensional projections of the score. For a random projection vector $v$ , the projected objective involves only a directional derivative:

v^{T} \nabla_{x} s (x; θ) v

which can be computed with a single vector-Jacobian product rather than the full Jacobian. Averaging over random projections recovers the full objective in expectation, at a fraction of the cost. In practice, denoising score matching is more common, but sliced score matching is the right tool when adding noise to the data is undesirable, for example with discrete or highly structured data where Gaussian perturbation is unnatural.

From a score to samples: Langevin dynamics

Learning a score is only half the story. To generate samples, we run Langevin dynamics, a stochastic iteration that climbs the score field while injecting noise so it does not collapse onto a single mode:

x_{t + 1} = x_{t} + \frac{ϵ}{2} s (x_{t}; θ) + ϵ z_{t}, z_{t} \sim N (0, I)

With a small step size $ϵ$ and enough steps, the iterates converge to samples from the model distribution. This is the sampling engine behind score-based generative models, and the ancestor of the reverse process in diffusion models, where the score is learned at many noise levels and the sampler walks from pure noise back to data.

The connection to diffusion models is direct. A diffusion model trained with denoising objectives is, in effect, learning a noise-conditional score function. The celebrated sample quality of diffusion models is therefore, in part, a vindication of score matching: learning gradients turned out to scale better than learning densities.

A concrete mental model

Consider a two-dimensional mixture of two Gaussians. The density has two peaks. The score field points towards the nearer peak everywhere except on the ridge between them, where it is near zero in the direction along the ridge.

A score network trained by denoising score matching will recover that field. Langevin sampling started from a random point will drift towards one peak, then jitter around it under the injected noise, producing samples from the mixture. It never needs to know the mixture weights as an explicit number: the weights are encoded implicitly in the shape of the field and in how much time the sampler spends near each peak.

Where this mental model helps is in diagnosing failure. If generated samples cluster on one mode, the learned score field is probably too steep there or too flat elsewhere. If samples look noisy, the sampler step size is probably too large for the learned field. Both are statements about a vector field, and both are easier to reason about than a statement about a normalising constant.

Limitations and common pitfalls

Score matching is elegant, but it has real failure modes that practitioners should know.

Low-density regions are poorly constrained. The objective is an expectation over data, so regions with no data contribute nothing. The learned score in empty regions is essentially arbitrary, and Langevin sampling can wander into those regions and get stuck. This is why annealed sampling across noise levels matters: large noise levels fill in the gaps and give the sampler a path back to the data.

The score does not determine the density uniquely in disconnected regions. If the data distribution has disconnected components, the score within each component says nothing about the relative mass of the components. Two distributions can share the same score field on the support of the data yet assign very different weights to each component. Score matching alone cannot fix mixture weights, and this shows up as mode-weight errors in practice.

Score estimation degrades in high dimensions with little data. The regression target in denoising score matching is noisy, and its variance grows with dimension and shrinks only slowly with the noise level. Small noise levels give a better approximation to the true score but a noisier target. Choosing a noise schedule, rather than a single noise level, is the standard compromise, and it is the reason modern score models are trained across a range of noise levels.

Langevin sampling is slow and sensitive to step size. A step size that is too large diverges; one that is too small mixes slowly. In practice, annealed Langevin dynamics with a decreasing step size schedule, or a proper diffusion sampler, is used instead of the basic iteration above.

Do not confuse the score with a classifier gradient. The score is the gradient with respect to the data, not with respect to parameters. A gradient of a classifier's output is a different object with a different meaning, and conflating the two leads to confused explanations of both score matching and classifier guidance.

When to use score matching

Score matching is a strong choice when the model class has an intractable normaliser, when sampling via a gradient-based iteration is acceptable, and when the data are continuous. It is the foundation of score-based generative models and diffusion models, and it is worth understanding even if you never train a score model directly, because so much of modern generative modelling is score estimation in disguise.

It is a poor choice when exact likelihoods are required, for example for model comparison by likelihood, anomaly detection by density threshold, or any setting where the normalising constant itself is the quantity of interest. In those settings, a normalising flow or an explicit likelihood model is more appropriate, and score matching should be used only as a diagnostic or a pretraining step.

Takeaway

The core idea fits in one sentence: the gradient of a log-density does not depend on the normalising constant, so matching gradients avoids the hardest part of density estimation. Hyvarinen's objective made that idea trainable, denoising score matching made it scalable, and Langevin dynamics made it generative. Every diffusion model you use today is, at heart, a score model that learned to point uphill.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial

Stein Variational Gradient Descent: Particles That Repel Their Way to the Posterior

Umesh Tharuka Malaviarachchi — Mon, 26 Oct 2026 10:47:46 +0000

Bayesian inference is simple to state and hard to compute. Given data, we want the posterior distribution over parameters. In almost every interesting model, that posterior has no closed form, its normalising constant is intractable, and direct sampling is expensive. Two classical families of tools attack the problem from opposite directions. Markov chain Monte Carlo (MCMC) gives asymptotically exact samples but can be slow to mix. Variational inference (VI) turns inference into optimisation, but usually forces the posterior into a chosen parametric family, such as a Gaussian, and then hopes that family is flexible enough.

Stein variational gradient descent (SVGD), introduced by Liu and Wang, offers a third way. It keeps a set of particles, moves them with a deterministic gradient-like rule, and lets the particles themselves represent the posterior. There is no parametric family to fit and no accept-or-reject step. The trick that makes it work is a repulsive force between particles. Without it, every particle would collapse to the same mode. With it, the particles spread out and trace the shape of the posterior.

The problem, stated precisely

Let $θ$ be a parameter vector and let $p (θ)$ be the target posterior, known only up to a constant. That is, we can evaluate $lo g p (θ)$ and its gradient, often called the score function, but we cannot evaluate the normalising constant and we cannot sample from $p$ directly.

Suppose we have $n$ particles $θ_{1}, \dots, θ_{n}$ drawn from some initial distribution $q$ . Our goal is to move these particles, step by step, so that their empirical distribution gets closer to $p$ . Closeness here is measured by Kullback-Leibler (KL) divergence:

KL (q ∥ p) = E_{θ \sim q} [lo g q (θ) - lo g p (θ)]

If we apply a smooth transformation $T (θ) = θ + ϵ ϕ (θ)$ to every particle, with a small step size $ϵ$ and a vector field $ϕ$ to be chosen, the distribution of the transformed particles changes. The question is: which $ϕ$ reduces the KL divergence fastest?

The Stein identity and the best direction

The answer comes from Stein's identity. For a smooth test function $ϕ$ that vanishes suitably at infinity, the following expectation under $p$ is zero:

E θ \sim p [\nabla θ lo g p (θ) \cdot ϕ (θ) + \nabla_{θ} \cdot ϕ (θ)] = 0

This identity is the reason SVGD needs only the score function $\nabla_{θ} lo g p (θ)$ , not the normalising constant. The constant disappears when we take the log gradient.

Stein variational gradient descent restricts $ϕ$ to lie in a reproducing kernel Hilbert space (RKHS) defined by a positive definite kernel $k$ . In that space, the direction of steepest descent of the KL divergence has a closed form. It is the expectation, over the current particle distribution, of a kernel-weighted score plus a kernel gradient term:

ϕ^{*} (θ) = E θ^{'} \sim q [k (θ^{'}, θ) \nabla θ^{'} lo g p (θ^{'}) + \nabla_{θ^{'}} k (θ^{'}, θ)]

That single equation contains the whole algorithm. Read its two terms separately, because they do opposite jobs.

The first term, $k (θ^{'}, θ) \nabla_{θ^{'}} lo g p (θ^{'})$ , is an attraction term. It is a kernel-weighted average of the score at the other particles. It pulls a particle uphill, toward regions where the posterior density is high, with nearby particles contributing more weight than distant ones.

The second term, $\nabla_{θ^{'}} k (θ^{'}, θ)$ , is a repulsion term. For the common radial basis function (RBF) kernel, this term points away from neighbouring particles. It stops particles from piling up on top of each other at a single mode. Without it, SVGD would be little more than parallel gradient ascent, and every particle would converge to the nearest mode. With it, the particles behave like a system of mutually repelling charges that still feels the pull of the posterior.

In practice we replace the expectation by the average over the finite particle set. The update for particle $i$ at iteration $t$ is:

θ_{i} \leftarrow θ_{i} + ϵ ϕ^^{(} θ_{i}), ϕ^^{(} θ_{i}) = \frac{1}{n} j = 1 \sum n [k (θ_{j}, θ_{i}) \nabla_{θ_{j}} lo g p (θ_{j}) + \nabla_{θ_{j}} k (θ_{j}, θ_{i})]

That is the entire algorithm. Initialise particles, repeatedly apply this update, and use the final particles as approximate posterior samples.

The kernel matters more than it looks

The usual choice is the RBF kernel:

k (θ, θ^{'}) = exp (- \frac{∣ θ - θ ^{'} ∣ ^{2}}{2 h ^{2}})

The bandwidth $h$ controls the range of interaction. If $h$ is too small, each particle only sees itself, the repulsion term vanishes, and the particles collapse. If $h$ is too large, the kernel is nearly constant, the repulsion is weak and indiscriminate, and the particles move almost in lockstep.

The standard practical fix is the median heuristic: set $h$ to the median pairwise distance between particles, recomputed at each iteration. This adapts the interaction range to the current spread of the particles, which is exactly what you want as the particles move from a diffuse initialisation to a concentrated posterior. It is a heuristic, not a theorem, but it works well in a wide range of problems and it removes one fragile tuning knob.

For the RBF kernel, the gradient term in the update has an explicit form:

\nabla_{θ_{j}} k (θ_{j}, θ_{i}) = - \frac{( θ _{j} - θ _{i} )}{h ^{2}} k (θ_{j}, θ_{i})

Notice the sign. When $θ_{j}$ is close to $θ_{i}$ , this term pushes $θ_{i}$ away from $θ_{j}$ , with a strength that decays with distance. That is the repulsion, written out in full.

A concrete example in code

Consider a simple target: a mixture of two Gaussians in two dimensions, with modes that are well separated. A single Gaussian variational approximation would be forced to pick one mode or sit uselessly between them. SVGD particles, by contrast, should split across both modes. The code below implements the update in NumPy for this toy target. It is short on purpose: the algorithm really is this short once the score function is available.

import numpy as np

def log_target_and_score(x):
    # Mixture of two 2D Gaussians, equal weights.
    # Returns log density (up to a constant) and its gradient.
    mu1 = np.array([2.0, 2.0])
    mu2 = np.array([-2.0, -2.0])
    d1 = x - mu1
    d2 = x - mu2
    log_p1 = -0.5 * np.sum(d1 ** 2, axis=-1)
    log_p2 = -0.5 * np.sum(d2 ** 2, axis=-1)
    m = np.maximum(log_p1, log_p2)
    log_p = m + np.log(0.5 * np.exp(log_p1 - m) + 0.5 * np.exp(log_p2 - m))
    w1 = np.exp(log_p1 - log_p)
    w2 = np.exp(log_p2 - log_p)
    score = -(w1[:, None] * d1 + w2[:, None] * d2)
    return log_p, score

def svgd(x, n_iter=500, eps=0.1):
    n = x.shape[0]
    for _ in range(n_iter):
        log_p, score = log_target_and_score(x)
        # Pairwise squared distances and the median heuristic.
        diff = x[:, None, :] - x[None, :, :]
        dist2 = np.sum(diff ** 2, axis=-1)
        h = np.median(np.sqrt(dist2 + 1e-8))
        h2 = h ** 2 + 1e-8
        K = np.exp(-dist2 / (2.0 * h2))
        # Kernel gradient term: -K * diff / h^2, summed over j.
        grad_K = -K[:, :, None] * diff / h2
        phi = (K @ score + np.sum(grad_K, axis=1)) / n
        x = x + eps * phi
    return x

rng = np.random.default_rng(0)
particles = rng.normal(size=(100, 2))
particles = svgd(particles)

Run this and plot the particles and you will see the intended behaviour: particles start as a single Gaussian blob, then split and settle around both modes, with their density roughly tracking the target. No particle is wasted sitting in the valley between the modes, because the score pulls it toward one mode or the other, and repulsion keeps the two groups from collapsing into points.

For a real Bayesian model, the only part that changes is log_target_and_score. For Bayesian logistic regression, for example, the score is the gradient of the log prior plus the log likelihood, which is available in closed form or by automatic differentiation. SVGD then gives you a particle approximation to the posterior over weights, from which you can compute predictive means, predictive variances, and credible intervals by averaging over particles.

Why SVGD is not just parallel MCMC or plain VI

It helps to place SVGD precisely relative to its neighbours.

Compared with MCMC, SVGD is deterministic given the initial particles. There is no random proposal and no rejection, so it can be much more sample efficient in the early phase, and it parallelises trivially across particles. The price is the loss of asymptotic exactness guarantees in the finite particle regime: with a fixed number of particles, SVGD converges to an approximation whose quality depends on the kernel, the bandwidth, and the particle count. MCMC, given infinite time, is exact. SVGD, given infinite particles, is also well understood in the mean-field limit, but you never have infinite particles.

Compared with parametric VI, SVGD is nonparametric in the sense that matters. The particle set can represent multiple modes, skewness, and other shapes that a Gaussian family cannot. The price is computational cost: each iteration costs $O (n^{2} d)$ for $n$ particles in $d$ dimensions, because every particle interacts with every other particle through the kernel. For a few hundred particles this is fine. For tens of thousands, it is not, and you need approximations such as mini-batch subsampling of particles or a structured kernel.

There is also a deeper connection worth knowing. In the limit of many particles, SVGD can be viewed as a gradient flow in the space of distributions, with the kernel defining the geometry of that flow. This is the same style of thinking that underlies Wasserstein gradient flows, except that the kernel geometry replaces the Wasserstein geometry. That viewpoint explains why the choice of kernel is not a minor implementation detail: it defines what "steepest descent" means in distribution space.

Limitations and common pitfalls

First, SVGD degrades in high dimensions. The RBF kernel suffers from the curse of dimensionality: in high dimensions, pairwise distances concentrate, the median heuristic becomes less informative, and the repulsion term weakens. Practical remedies include using a kernel on a lower dimensional representation, using an anisotropic kernel with a learned length scale per dimension, or combining SVGD with a dimension reduction step. Do not expect out-of-the-box SVGD with a single isotropic RBF kernel to work well on a thousand dimensional posterior without such adjustments.

Second, too few particles is a silent failure mode. With, say, 20 particles on a multimodal target, SVGD may still look plausible: the particles spread out and find high density regions. But the particle distribution is a coarse approximation, and estimates of posterior variance or tail probabilities will be unreliable. The diagnostic is simple: increase the particle count and check whether your quantity of interest changes. If it moves materially, you did not have enough particles.

Third, the step size needs care. A fixed step size that is too large causes oscillation or divergence, especially early on when gradients are large. In practice, practitioners often use an adaptive optimiser such as AdaGrad or Adam on the SVGD update direction, treating $ϕ^^{*}$ as a gradient. This is a heuristic layered on top of the theory, but it stabilises training considerably. Monitor the movement of particles, not just the objective, because SVGD has no simple scalar loss that tells you it has converged.

Fourth, SVGD needs a differentiable log density. If your model has discrete parameters, or if the score is unavailable or noisy, the basic algorithm does not apply directly. Variants exist for discrete and gradient-free settings, but they change the theory, so do not force the standard update onto a problem where the score does not exist.

Fifth, do not confuse SVGD particles with independent posterior samples. They are a deterministic, interacting approximation. For predictive averaging this distinction rarely matters much, but for estimating Monte Carlo error bars it does. If you need honest uncertainty about your Monte Carlo estimate itself, SVGD particles alone do not give it to you.

When to reach for SVGD

SVGD earns its place when three conditions hold: you can compute the score function, the dimension is moderate, and the posterior may be multimodal or otherwise non-Gaussian in a way that matters for your decision. Bayesian neural networks with a modest number of parameters, Bayesian logistic regression with structured priors, and probabilistic models where predictive uncertainty drives an action are all natural fits.

It is less compelling when you need asymptotically exact samples and can afford a long MCMC run, when the model is high dimensional and a well tuned parametric VI approximation is already good enough, or when the score function is unavailable. As with most inference methods, the right question is not whether SVGD is universally better, but whether its particular trade-off, flexibility for cost, matches your problem.

The mental image to keep is the one from the update equation: each particle is pulled toward high posterior density and pushed away from its neighbours. Attraction alone gives you optimisation. Repulsion alone gives you diversity. Together, and only together, they give you inference.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial

Optimal Transport: Why Wasserstein Distance Sees Geometry That KL Divergence Misses

Umesh Tharuka Malaviarachchi — Mon, 26 Oct 2026 07:47:33 +0000

Optimal Transport: Why Wasserstein Distance Sees Geometry That KL Divergence Misses

Two probability distributions can be very different in a way that KL divergence barely notices, or notices too harshly. Move a narrow bump of probability a little to the right, and KL divergence changes smoothly. Move it far away so the two bumps no longer overlap, and KL divergence can blow up or become undefined, even though the intuitive answer is simple: the distance is just how far the mass had to travel.

Optimal transport makes that intuition precise. It asks: what is the cheapest way to reshape one distribution into another, when moving a unit of mass across a distance incurs a cost? The answer, the Wasserstein distance, measures difference in the geometry of the underlying space, not only overlap of densities. That one change explains why optimal transport shows up in generative modelling, domain adaptation, shape comparison, and distribution matching, and also why it can be expensive and fragile in high dimensions.

The earth mover intuition

Imagine each distribution as a pile of earth spread over a landscape. To turn the first pile into the second, you may dig earth from some locations and deposit it elsewhere. Moving earth a short distance is cheap. Moving it across the landscape is expensive. The total cost depends on both how much earth moves and how far it moves.

KL divergence asks a different question. It compares densities point by point: at each location, how surprising is one density under the other? If the two distributions put mass in disjoint places, that pointwise comparison breaks down. Optimal transport instead compares mass, and allows mass to travel.

A simple example makes the contrast sharp. Let P put all its mass at 0 and let Q put all its mass at t, where t is a positive number. Intuitively, the distance between P and Q should grow with t. The Wasserstein distance does exactly that: it equals t. KL divergence, by contrast, is not a useful guide here, because the supports do not overlap. This is not a pathological curiosity. Real data often live on low dimensional manifolds, and two empirical distributions rarely share support in a meaningful way.

From Monge to Kantorovich

The classical formulation is due to Monge. Given a source distribution and a target distribution, find a map T that pushes the source onto the target while minimising total transport cost. The difficulty is that such a map may not exist: a single source point cannot split its mass between two target points under a deterministic map.

Kantorovich relaxed the problem by allowing mass to split. Instead of a map, we look for a coupling, a joint distribution whose marginals are the two distributions of interest. The coupling says how much mass travels from each source location to each target location. The optimal transport cost is the minimum, over all couplings, of the expected transport cost.

For distributions P and Q on a space with a ground cost c, the Kantorovich problem is:

π \in Π (P, Q) in f \int c (x, y) d π (x, y)

Here Pi(P, Q) is the set of all couplings of P and Q. When the ground cost is a power of a metric, for example c(x, y) = ||x - y||^p with p at least 1, the p-th root of this minimum is the p-Wasserstein distance:

W_{p} (P, Q) = (π \in Π (P, Q) in f \int ∣∣ x - y ∣ ∣^{p} d π (x, y))^{1/ p}

The two most used cases are p = 1 and p = 2. W1 is often called the earth mover distance. W2 has extra structure: under regularity conditions, the optimal map is the gradient of a convex function, a result known as Brenier's theorem. That structure is why W2 appears so often in theory and in algorithms that need a genuine map rather than only a distance.

A useful way to read W1 is through its dual form. For distributions with finite first moments:

W_{1} (P, Q) = ∣∣ f ∣∣ L \leq 1 sup (E X \sim P [f (X)] - E_{Y \sim Q} [f (Y)])

The supremum is over all 1-Lipschitz functions f. In words, W1 is the largest difference in expectation that a function with slope at most 1 can detect. This dual view is the bridge to Wasserstein GANs: a critic network, constrained to be approximately Lipschitz, tries to maximise exactly this gap, while the generator tries to minimise it.

The discrete case: a linear program

In practice we usually have samples, not densities. Suppose P is supported on points x1 to xm with weights a1 to am, and Q is supported on points y1 to yn with weights b1 to bn. A transport plan is an m by n matrix Gamma with nonnegative entries, where row i sums to ai and column j sums to bj. Entry Gamma(i, j) is the mass moved from xi to yj. If C(i, j) is the cost of moving mass from xi to yj, the discrete problem is:

Γ \geq 0 min i = 1 \sum m j = 1 \sum n Γ_{i, j} C_{i, j}

subject to:

j = 1 \sum n Γ_{i, j} = a_{i}, i = 1 \sum m Γ_{i, j} = b_{j}

This is a linear program. That is good news for correctness and bad news for scale: a naive solution costs far more than most machine learning routines can afford when m and n are large. Much of modern optimal transport research is about approximations that keep the geometry while making this problem tractable.

One dimension is special

In one dimension, optimal transport has a closed form that is worth knowing, because it also gives a fast sanity check in higher dimensional work. For distributions on the real line with cumulative distribution functions F_P and F_Q, and quantile functions F_P^{-1} and F_Q^{-1}, the p-Wasserstein distance is:

W_{p} (P, Q) = (\int_{01} ∣ F_{P - 1} (u) - F_{Q - 1} (u) ∣^{p} d u)^{1/ p}

For p = 1, this simplifies to the area between the two cumulative distribution functions:

W_{1} (P, Q) = \int_{- \infty \infty} ∣ F_{P} (x) - F_{Q} (x) ∣ d x

The intuition is sorting. In one dimension, the optimal plan never crosses itself: the smallest source mass goes to the smallest target mass, the second smallest to the second smallest, and so on. That monotone matching is both optimal and easy to compute.

The code below computes W1 between two equal sized samples in one dimension using exactly this sorted matching. It needs no solver and no density estimate.

import numpy as np

def wasserstein_1d_sample(x, y):
    x = np.sort(np.asarray(x, dtype=float))
    y = np.sort(np.asarray(y, dtype=float))
    if x.shape != y.shape:
        raise ValueError("This shortcut needs equal sized samples.")
    return float(np.mean(np.abs(x - y)))

# Two samples from similar but shifted distributions
rng = np.random.default_rng(0)
source = rng.normal(loc=0.0, scale=1.0, size=1000)
target = rng.normal(loc=2.0, scale=1.0, size=1000)
print(wasserstein_1d_sample(source, target))

The output is close to 2, the size of the shift, as intuition suggests. If the two samples came from the same distribution, the value would be small but not zero, because finite samples differ by chance. That sampling noise matters: always compare an observed Wasserstein distance against a null value obtained by splitting one sample in two, or by permuting labels, before claiming that two distributions differ.

Entropic regularisation and the Sinkhorn algorithm

Exact optimal transport is expensive, and its solution can be sparse and unstable: a tiny change in the data can switch the optimal plan. Entropic regularisation adds a penalty that favours diffuse, smooth plans. The regularised problem is:

Γ \in Π (a, b) min i, j \sum Γ_{i, j} C_{i, j} + ε i, j \sum Γ_{i, j} lo g Γ_{i, j}

Here epsilon is a positive regularisation parameter. Larger epsilon gives a smoother plan and a faster algorithm, but a more biased distance. Smaller epsilon approaches exact transport but needs more iterations and more careful numerics.

The regularised problem can be solved with Sinkhorn iterations, which alternate between rescaling rows and columns of a kernel matrix derived from the cost matrix. Each iteration is a matrix-vector product, so the method is fast on GPUs and differentiable, which is why it appears inside neural networks and in differentiable sorting, matching, and clustering layers.

The trade-off must be stated plainly. Sinkhorn distance is not the Wasserstein distance. It is a regularised approximation, and its value depends on epsilon. Reporting a Sinkhorn number without epsilon, and without the ground cost, is like reporting a loss without saying which loss. If the downstream decision depends on the exact value, validate against exact transport on a subsample.

Where Wasserstein helps, and where it misleads

Wasserstein distance earns its place when geometry matters. Three settings are typical.

First, generative modelling. A generator that produces samples slightly off the data manifold should be penalised in proportion to how far off it is, not infinitely because densities do not overlap. The WGAN idea uses the W1 dual to build a critic with that property. In practice the Lipschitz constraint is enforced only approximately, by weight clipping or gradient penalty, so a WGAN loss is an approximation to an approximation. Treat it as a training signal, not as a certified distance.

Second, domain adaptation and distribution matching. If source and target domains differ by a shift, rotation, or deformation of features, a transport plan can align them in a way that a pointwise density ratio cannot. The plan itself is often more useful than the scalar distance: it tells you which source examples correspond to which target examples.

Third, comparing structured objects. Distributions over images, point clouds, word embeddings, or histograms can be compared with a ground cost that respects the structure of the objects, for example Euclidean distance between points in a cloud, or cosine distance between embeddings. The choice of ground cost is the modelling choice. Wasserstein distance inherits its meaning from that cost.

The limitations are equally important.

Sample complexity in high dimensions is the main one. Estimating a Wasserstein distance from samples gets harder as dimension grows, roughly because the space has more directions in which two samples can differ. In high dimensions, use sliced Wasserstein distances, which average one dimensional Wasserstein distances over random projections, or use embeddings and dimensionality reduction first, and report that choice.

The ground metric can dominate the result. Change the feature scaling and you change the transport cost, even if the data are unchanged. Standardise features with a reason, not by habit, and check sensitivity to that scaling before drawing conclusions.

Mass must be conserved in the standard formulation. If one distribution has extra mass that the other does not, for example outliers or a new class, standard transport is forced to move that mass somewhere, possibly at high cost. Unbalanced optimal transport relaxes exact marginal constraints and is the right tool when creation or destruction of mass is part of the problem.

Finally, a small Wasserstein distance does not mean two distributions are the same in every respect. It means they are close under one chosen ground cost. Two distributions can have the same mean, the same Wasserstein distance to a reference, and very different tails, modes, or conditional structure. Use Wasserstein as one diagnostic among several, alongside moment checks, conditional checks, and task performance.

A practical checklist

Before using a Wasserstein distance in a project, answer these questions in order:

What is the ground cost, and why does it reflect what "far" means in this problem?
Are the features scaled so that the ground cost is meaningful, or is one feature silently dominating?
Is exact transport feasible, or is entropic regularisation needed, and if so, what is epsilon?
What is the null value of the distance between two samples from the same distribution at this sample size?
Does mass need to be conserved, or is unbalanced transport more honest?
Is the transport plan itself useful, or only the scalar distance?

If any answer is unclear, the number that comes out of the solver will be precise and uninformative. Optimal transport is powerful precisely because it forces you to say what distance means before you measure it.

Takeaway

KL divergence compares densities at the same point. Wasserstein distance asks what it costs to move mass from one distribution to another. When distributions are separated, live on different manifolds, or differ by a geometric shift, that transport view gives a smooth, interpretable answer where a pointwise comparison fails. The price is computation, sample complexity in high dimensions, and a strong dependence on the ground cost. Choose the ground cost carefully, report regularisation honestly, and use the transport plan, not only the distance, whenever the correspondence between samples is what you actually need.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial

The Information Bottleneck: Why the Best Representations Forget Most of What They See

Umesh Tharuka Malaviarachchi — Mon, 26 Oct 2026 05:02:48 +0000

A neural network trained to classify images sees millions of pixels. A good representation keeps almost none of them. It keeps the shape of a cat, the texture of fur, the arrangement of ears, and throws away the exact shade of the background, the sensor noise, and the compression artefacts. This is not a failure of the model. It is the whole point of learning a representation.

The Information Bottleneck (IB), introduced by Tishby, Pereira, and Bialek in 1999, turns this intuition into a precise optimisation problem: find a representation that is as compressed as possible while still preserving everything relevant about a target variable. It gives us a language for asking what a model should remember, what it should forget, and how to trade one against the other.

The setup

Suppose we have an input variable X and a target variable Y, drawn from some joint distribution p(x, y). Think of X as an image and Y as its label. We want a representation Z, typically a stochastic function of X, that captures the information in X that matters for Y, and nothing else.

Two quantities matter:

Compression: how much information Z retains about X, measured by the mutual information I(X; Z). Smaller is more compressed.
Prediction: how much information Z retains about Y, measured by I(Z; Y). Larger is more predictive.

The Information Bottleneck formalises the trade-off as a constrained optimisation: minimise I(X; Z) subject to I(Z; Y) being at least some value, or equivalently, minimise the Lagrangian

L [p (z ∣ x)] = I (X; Z) - β I (Z; Y)

where $β \geq 0$ is a Lagrange multiplier that controls the trade-off. Small $β$ means aggressive compression, even at the cost of predictive power. Large $β$ means prediction dominates, and the representation is allowed to carry more of X.

This is deliberately different from ordinary supervised learning. A classifier trained with cross-entropy alone has no explicit pressure to forget anything. The bottleneck adds that pressure, and that pressure is where generalisation comes from: a representation that carries only task-relevant information cannot overfit to irrelevant details, because it never stored them.

Why mutual information, and not something simpler

It is tempting to ask why we need mutual information at all. Why not just penalise the dimensionality of Z, or the norm of the representation?

The answer is that mutual information is invariant to invertible transformations of Z. A representation can be rescaled, rotated, or reparametrised without changing what it knows. Dimensionality and norm penalties punish the encoding, not the information content. Mutual information punishes the information content itself, which is what we actually care about.

Mutual information is also symmetric and distribution-free in a useful sense. Recall its definition:

I (X; Z) = x, z \sum p (x, z) lo g \frac{p ( x , z )}{p ( x ) p ( z )}

It measures how far the joint distribution is from independence, in nats or bits depending on the logarithm. When Z is independent of X, the ratio is 1, the logarithm is 0, and I(X; Z) = 0. When Z determines X exactly, I(X; Z) equals the full entropy H(X).

The catch, and it is a serious one, is that mutual information between high-dimensional continuous variables is notoriously hard to estimate. This is why the original IB remained mostly a theoretical tool for years, and why the practical versions of it look quite different from the clean Lagrangian above.

The self-consistent solution

For discrete variables, the IB problem has a beautiful self-consistent solution, analogous to the Blahut-Arimoto algorithm for rate-distortion theory. The optimal encoder takes the form

p (z ∣ x) = \frac{p ( z )}{Z ( x , β )} exp (- β D_{K L} [p (y ∣ x) ∣ p (y ∣ z)])

where $Z (x, β)$ is a normalising constant. Read this carefully, because it is the heart of the theory: an input x is assigned to a representation z with probability that decays exponentially in the KL divergence between what x predicts about y and what z predicts about y. Inputs that imply similar predictions about Y get mapped to similar representations. Inputs that imply different predictions get separated.

This is clustering, but with a very specific notion of similarity: two inputs are similar if they carry the same information about the target. The bottleneck does not cluster by how inputs look, but by what they mean for the task.

Varying $β$ traces out a curve in the information plane, with I(X; Z) on the horizontal axis and I(Z; Y) on the vertical axis. At $β = 0$ , the solution is trivial: Z is independent of everything, and both informations are 0. As $β$ grows, the curve rises, eventually saturating at I(Z; Y) = I(X; Y), the total relevant information available. The curve is concave, and its slope at any point is $1/ β$ . Points on this curve are optimal trade-offs. Points below it are achievable but wasteful. Points above it are impossible.

For jointly Gaussian X and Y, the solution is analytic: the optimal Z is a noisy linear projection of X, and the trade-off curve can be written in closed form in terms of the canonical correlations between X and Y. This Gaussian case is worth knowing because it connects the bottleneck to canonical correlation analysis and to linear dimensionality reduction, and it gives a sanity check for any approximate algorithm.

The information plane and deep learning

The IB became famous in deep learning through a 2017 paper by Shwartz-Ziv and Tishby, which claimed that stochastic gradient descent training of deep networks proceeds in two phases: a fast fitting phase, in which I(Z; Y) rises, followed by a slow compression phase, in which I(X; Z) falls while I(Z; Y) stays roughly constant. The claim was that this compression phase is what makes deep networks generalise, and that it is driven by the noise in SGD.

This claim sparked one of the liveliest debates in recent machine learning. Follow-up work showed that the compression phase depends heavily on the activation function (it appears clearly with saturating activations like tanh, and much less with ReLU), on the estimator used for mutual information, and on whether the network is deterministic or stochastic. Later studies argued that the observed compression was partly an artefact of the binning estimator used to measure mutual information in deterministic networks, where the true mutual information is actually infinite or vacuous.

The debate is instructive rather than embarrassing. It forced the field to be precise about what it means for a deterministic network to compress information at all, and it pushed the practical work toward stochastic networks, where mutual information is well-defined. Whether or not every deep network literally traverses the information plane, the IB remains the cleanest framework we have for stating what a good representation should do.

The Variational Information Bottleneck

The practical breakthrough came with the Variational Information Bottleneck (VIB), introduced by Alemi and co-authors in 2017. The idea is to replace the intractable mutual information terms with variational bounds that can be optimised with stochastic gradient descent.

The compression term I(X; Z) is upper-bounded using a variational marginal r(z):

I (X; Z) \leq E x [D K L (p (z ∣ x) ∣ r (z))]

The prediction term I(Z; Y) is lower-bounded using a variational decoder q(y|z):

I (Z; Y) \geq E x, y E z \sim p (z ∣ x) [lo g q (y ∣ z)] + H (Y)

Dropping the constant H(Y), the VIB objective becomes

L V I B = E x, y E z \sim p (z ∣ x) [- lo g q (y ∣ z)] + β E x [D_{K L} (p (z ∣ x) ∣ r (z))]

This is immediately recognisable: the first term is a cross-entropy loss, and the second is a KL regulariser pulling the encoder toward a prior. If you have trained a variational autoencoder, you have optimised something very close to this. The difference is the target: a VAE reconstructs X, while the VIB predicts Y.

Here is a minimal PyTorch sketch of the idea, with a Gaussian encoder and a standard normal prior:

import torch
import torch.nn as nn

class VIBEncoder(nn.Module):
    def __init__(self, input_dim, z_dim):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(input_dim, 256), nn.ReLU())
        self.mu = nn.Linear(256, z_dim)
        self.logvar = nn.Linear(256, z_dim)

    def forward(self, x):
        h = self.net(x)
        mu, logvar = self.mu(h), self.logvar(h)
        std = torch.exp(0.5 * logvar)
        z = mu + std * torch.randn_like(std)  # reparametrisation trick
        kl = -0.5 * torch.sum(1 + logvar - mu**2 - logvar.exp(), dim=1)
        return z, kl

# Training step: cross-entropy on the prediction plus beta times the KL term
z, kl = encoder(x)
logits = classifier(z)
loss = nn.functional.cross_entropy(logits, y) + beta * kl.mean()
loss.backward()

Two practical details matter here. First, the encoder must be stochastic: without the sampling step, the KL term is meaningless and the bottleneck does nothing. Second, $β$ is a genuine hyperparameter that needs tuning, not a value to set and forget. Too small, and the KL term is ignored. Too large, and the representation collapses to the prior and predicts nothing. A common practice is to anneal $β$ from a small value upward during training.

What the bottleneck buys you

The VIB and its relatives have delivered real, measurable benefits in several settings:

Robustness. Representations trained with a bottleneck are less sensitive to input perturbations, including adversarial ones, because irrelevant details were never encoded in the first place. An adversary cannot exploit information the model does not have.
Out-of-distribution detection. The KL term gives a natural signal: inputs far from the training distribution tend to produce encodings far from the prior, which can be used as an uncertainty or anomaly score.
Disentanglement and interpretability. By limiting the information budget, the bottleneck forces the representation to spend its capacity on the most predictive factors, which often correspond to interpretable structure.
Generalisation bounds. The IB connects directly to generalisation theory: bounds on the generalisation gap can be expressed in terms of I(X; Z), so compressing the representation provably limits overfitting in a precise, distribution-dependent sense.

None of this is magic, and none of it is free. The bottleneck trades accuracy on the training distribution for robustness elsewhere, and the trade is only worth making when the irrelevant information in X is genuinely harmful rather than merely unused.

Limitations and common pitfalls

Mutual information estimation is hard. The variational bounds are bounds, not equalities, and a loose bound means you are optimising a different objective than the one you wrote down. If your variational marginal r(z) is too rigid, the compression penalty is mismeasured.
The bottleneck can discard information you later need. A representation optimised for one target Y may be useless for a related target Y'. This is the transfer learning trap: aggressive compression is task-specific, and a representation that forgot everything except what one task needed cannot be reused for another.
Deterministic networks and the IB do not mix cleanly. For a deterministic function Z = f(X), the mutual information I(X; Z) is infinite for continuous X (or equal to H(X) for discrete X), so the Lagrangian is vacuous. Claims about compression in deterministic networks need a different formalism, such as the recent work on stochastic relaxations or on geometric notions of compression.
Beta is not a free lunch. There is no universal value of $β$ that works across tasks. It must be tuned, and the information plane for your problem should be inspected, not assumed.
Do not confuse the bottleneck with ordinary regularisation. Weight decay and dropout also limit capacity, but they do so in parameter space, not information space. The bottleneck is a statement about what the representation knows, which is a stronger and more interpretable claim.

The takeaway

The Information Bottleneck asks the right question about representation learning: not "how much can this model remember?" but "how little does it need to remember to do the job?" The answer is a curve, not a number, and the interesting work is choosing where on that curve to sit.

When you train your next model, ask what it is carrying that it does not need. The chances are that a large fraction of its representation is memorised noise, and that the model would be simpler, more robust, and easier to understand if it forgot most of what it saw.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial

Conformal Prediction: Valid Uncertainty Without Distributional Assumptions

Umesh Tharuka Malaviarachchi — Mon, 26 Oct 2026 01:59:41 +0000

Most machine learning models give you a point prediction and stop there. A house price, a risk score, a class label. The number looks precise, but precision is not the same as reliability. If a model says a patient will recover in five days, a decision maker does not only need the five. They need to know how wrong that number could plausibly be, and how often the stated range can be trusted.

Conformal prediction answers that question with unusual honesty. It turns almost any predictor into a prediction set or interval with a finite-sample coverage guarantee, and it does so without assuming that residuals are Gaussian, that the model is correctly specified, or that you know the true data generating process. The price is modest: you need a held-out calibration set, and you need exchangeability. The reward is a guarantee you can state plainly.

What "valid" means here

Suppose you choose a target miscoverage level, say $α = 0.1$ . A conformal prediction set $C (X)$ aims to satisfy:

P (Y_{n + 1} \in C (X_{n + 1})) \geq 1 - α

That is marginal coverage. Over many fresh test points drawn in the same way as the calibration data, at least 90 percent of the true outcomes should fall inside their sets. It is not a promise for one specific input, and it is not a promise for every subgroup. That distinction matters, and we will return to it.

The guarantee is finite-sample, not asymptotic. It does not only hold when the calibration set grows infinitely large. It holds for the calibration set you actually have, provided the exchangeability assumption holds. Exchangeability is weaker than independence plus identical distribution in the formal sense, but for practical purposes you can read it as: the calibration points and the future test point are drawn from the same process, with no ordering, drift, or selection effect that makes the test point systematically different.

That is the whole trick, and also the whole vulnerability. Conformal prediction does not fix distribution shift. It converts a calibration sample into a calibrated statement about a future sample from the same distribution.

Split conformal prediction in four steps

Split conformal prediction is the simplest and most widely used variant. You split your data into a training set and a calibration set, fit the model once, and never refit it during calibration.

Step 1: Fit a predictor $f^$ on the training set. It can be a linear model, a gradient boosting model, a neural network, or anything else. Conformal prediction is model agnostic.

Step 2: Choose a nonconformity score. This score measures how unusual an observed outcome is, given the prediction. For regression, the classic choice is the absolute residual:

s_{i} = ∣ Y_{i} - f^(X_{i}) ∣

A large score means the model was badly wrong on that calibration point. A small score means it was close.

Step 3: Compute the scores on the calibration set, then take a suitable quantile. If you have $n$ calibration points and want coverage $1 - α$ , you do not take the plain $1 - α$ quantile. You take the adjusted level:

q^= Quantile (s_{1}, \dots, s_{n}; \frac{⌈( n + 1 ) ( 1 - α )⌉}{n})

The ceiling and the $n + 1$ correction are not pedantry. They are exactly what turns an approximate empirical quantile into a finite-sample guarantee. Intuitively, the new test point is equally likely to rank anywhere among the $n + 1$ points (the calibration points plus itself) under exchangeability, so you must leave room for that extra rank.

Step 4: For a new input $X_{n + 1}$ , return every outcome whose score would be at most $q^$ . With the absolute residual score, that gives the familiar symmetric interval:

C (X_{n + 1}) = [f^(X_{n + 1}) - q^, f^(X_{n + 1}) + q^]

That is it. No distributional assumption about the residuals was used. No second model was fitted. The calibration scores did all the work.

Why the guarantee works, in one paragraph

Under exchangeability, the scores of the calibration points and the future point are exchangeable too. Rank the $n + 1$ scores. The future score is equally likely to occupy any rank. The adjusted quantile is chosen so that the future score falls at or below it with probability at least $1 - α$ . Falling at or below it is precisely the event that the true outcome lies inside the prediction set. This rank argument is why the method is so general: it never asks what distribution produced the scores, only that their ordering is not rigged by the split.

A worked example in code

The code below implements split conformal prediction for regression with NumPy. It is deliberately minimal so the mechanics are visible: fit, score, take the adjusted quantile, build intervals.

import numpy as np
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.model_selection import train_test_split

# X, y are your features and targets. Split into train, calibration, and test.
X_train, X_rest, y_train, y_rest = train_test_split(X, y, test_size=0.4, random_state=0)
X_cal, X_test, y_cal, y_test = train_test_split(X_rest, y_rest, test_size=0.5, random_state=1)

model = HistGradientBoostingRegressor(random_state=0)
model.fit(X_train, y_train)

# Step 2 and 3: absolute residual scores, then the adjusted quantile.
cal_pred = model.predict(X_cal)
scores = np.abs(y_cal - cal_pred)

alpha = 0.1
n = len(scores)
level = np.ceil((n + 1) * (1 - alpha)) / n
q_hat = np.quantile(scores, level, method="higher")

# Step 4: intervals for new points.
test_pred = model.predict(X_test)
lower = test_pred - q_hat
upper = test_pred + q_hat

coverage = np.mean((y_test >= lower) & (y_test <= upper))
width = np.mean(upper - lower)
print(f"Empirical coverage: {coverage:.3f} (target at least {1 - alpha})")
print(f"Mean interval width: {width:.3f}")

Two details deserve attention. First, the calibration set must not have been used to fit the model, or the scores will be optimistically small and the intervals too narrow. Second, the quantile method matters at small $n$ . The "higher" method used above is a safe choice for the finite-sample argument; other interpolation rules can break the exact guarantee even if they look similar.

If your calibration set is small, the intervals will still be valid on average, but their width will be noisy. Validity is about coverage over repeated draws of calibration and test data, not about one lucky or unlucky calibration sample giving a stable width. In practice, a few hundred calibration points are a reasonable starting point for a 90 percent target, and a thousand or more gives much more stable widths.

Constant width is a limitation, not a feature

The absolute residual version gives every test point an interval of the same width: $2 q^$ . That is honest on average, but it ignores heteroscedasticity. A region where the model is accurate gets the same width as a region where it is erratic. For many applications, that is wasteful and hard to explain.

A natural fix is a normalised score. Suppose you also have a difficulty estimate $σ^(X)$ , for example from a model trained to predict absolute residuals, or from the spread of an ensemble. Then define:

s_{i} = \frac{∣ Y _{i} - f ^ ( X _{i} ) ∣}{σ ^ ( X _{i} )}

The prediction interval becomes:

C (X_{n + 1}) = [f^(X_{n + 1}) - q^σ^(X_{n + 1}), f^(X_{n + 1}) + q^σ^(X_{n + 1})]

The coverage guarantee is unchanged, because the rank argument does not care which score you chose. What changes is adaptivity: intervals widen where the model is expected to struggle and narrow where it is expected to be sharp. The catch is that $σ^$ must be fitted on training data, not on the calibration set, or you again leak calibration information into the scores.

Conformalised quantile regression goes one step further. Instead of starting from a point predictor, you fit lower and upper quantile models, say at levels $α /2$ and $1 - α /2$ , then conformalise their residuals. The score becomes the distance outside the initial quantile interval, and the final interval corrects that interval just enough to restore coverage. When the quantile models capture heteroscedasticity well, this often gives sharper, more adaptive intervals than a normalised residual score.

The lesson generalises: conformal prediction is a wrapper, not a model. A better base model or a better score usually gives narrower sets at the same coverage. The wrapper guarantees coverage; the score decides efficiency.

Classification: prediction sets, not just labels

In classification, the output is a set of plausible labels rather than an interval. A simple score is one minus the predicted probability of the true class:

s_{i} = 1 - p^_{Y_{i}} (X_{i})

At test time, you include every class whose score is at most $q^$ , equivalently every class with predicted probability at least $1 - q^$ . Easy inputs get singleton sets. Hard or ambiguous inputs get larger sets. That behaviour is exactly what you want from an abstention or triage system: the size of the set is itself a signal.

More refined scores, such as adaptive prediction sets, accumulate sorted class probabilities until the true class is reached, with a randomisation step to avoid over-conservatism from discrete scores. The details matter for efficiency, but the principle is the same: define a score, calibrate its quantile, include everything that scores well enough.

A common mistake is to threshold softmax probabilities directly, for example by keeping classes above 0.5, and to call the result calibrated. Softmax probabilities from a modern neural network are often poorly calibrated, and an arbitrary threshold has no coverage guarantee. Conformal prediction replaces the arbitrary threshold with a quantile chosen from held-out data.

What conformal prediction does not promise

Marginal coverage is weaker than most people first assume. A method with 90 percent marginal coverage can still have 70 percent coverage for one subgroup and 98 percent for another, as long as the average works out. If you need coverage within each group, you need conditional or group-conditional methods. Mondrian conformal prediction is the classic fix: split the calibration set by group and calibrate a separate quantile per group. You pay for that guarantee with smaller calibration samples per group and wider intervals for small groups, which is an honest reflection of how much you actually know about them.

Full conditional coverage, meaning 90 percent coverage for every single $x$ , is impossible in general without further assumptions. This is a known impossibility result, not a gap in implementation. Any method claiming exact conditional coverage for arbitrary distributions should be treated with scepticism.

Distribution shift is the other hard boundary. If the deployment distribution differs from the calibration distribution, the guarantee does not transfer. Weighted conformal prediction can handle known or estimable shifts, such as covariate shift with a known likelihood ratio, but it cannot conjure validity from a calibration set that no longer resembles deployment. In practice, this means you should monitor coverage on fresh labelled data after deployment, refresh the calibration set when the distribution moves, and treat a conformal interval built on stale data as a historical artefact, not a live guarantee.

Finally, conformal prediction does not make a bad model good. If the base predictor is poor, the intervals will be wide, because the calibration scores will be large. Wide intervals are the method telling the truth about the model. Shrinking them by hand, by tuning $α$ after looking at test results, or by reusing calibration data for training, destroys the guarantee and replaces it with wishful thinking.

Conformal prediction versus Bayesian and ensemble uncertainty

Bayesian methods give posterior predictive intervals, but only under the prior and likelihood you chose. If those are misspecified, the intervals can be systematically wrong, and there is no finite-sample correction waiting to rescue them. Ensembles and dropout-based uncertainty give useful heuristics, but their spread is not calibrated by construction either.

Conformal prediction sits at a different level. It does not compete with a Bayesian model or an ensemble; it can wrap either of them. Use the Bayesian posterior mean as $f^$ and the posterior standard deviation as $σ^$ , then conformalise. You keep the structure of the Bayesian model and gain a distribution-free coverage statement on top. That combination is often the most defensible choice in a high-stakes setting: a sensible model, plus a wrapper that does not trust the model.

Practical checklist

Before you ship conformal intervals, check these points:

The calibration set was not used for training, tuning, or early stopping. Any reuse leaks optimism into the scores.
The calibration set is large enough that the quantile is stable, and the split was random with respect to the deployment process.
The score matches the task: absolute residuals for a first pass, normalised or quantile-based scores when heteroscedasticity matters, set-based scores for classification.
You report both coverage and mean width. Coverage alone is easy to game by returning huge sets. Width alone is meaningless without coverage.
You monitor coverage after deployment and have a plan to recalibrate when drift appears.
You state the guarantee correctly: marginal coverage at the chosen level, under exchangeability, for the calibration distribution. No more, no less.

Takeaway

Conformal prediction is rare in machine learning: a method whose main claim is modest, precise, and actually true. It does not promise that your model is right. It promises that, if tomorrow looks like your calibration data, your stated uncertainty will be right at least as often as you claimed. For decisions where overconfidence is expensive, that is a much more useful promise than a sharper-looking point prediction.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial

Spectral Bias: Why Neural Networks Learn Low Frequencies First

Umesh Tharuka Malaviarachchi — Sun, 25 Oct 2026 13:55:24 +0000

Train a neural network on a wiggly target function and watch what happens, epoch by epoch. The network does not learn the whole function at once. It first fits a smooth, coarse approximation, and only later starts to capture the sharp wiggles, spikes, and fine detail. This is not a bug in your optimiser or a quirk of one architecture. It is a systematic preference known as spectral bias, also called the frequency principle.

Understanding spectral bias changes how you think about training, generalisation, and failure modes. It explains why networks generalise surprisingly well on smooth data, why they struggle with high-frequency detail, and why some tricks that look like regularisation are really just frequency control in disguise.

The intuition in one picture

Imagine trying to sketch a coastline. You would start with the broad shape of the coast, then add bays, then inlets, then rocky detail. A neural network does something similar, but in frequency space rather than in geographic space.

Low frequencies describe broad, slow variation: the overall trend, the big humps, the global shape. High frequencies describe rapid variation: sharp edges, oscillations, and fine texture. During gradient descent, the low-frequency components of the target function are fitted much faster than the high-frequency components.

That ordering matters. If you stop training early, you get a smooth model. If you train for a very long time, you eventually fit the detail too, including the noise. Early stopping, in this light, is not just a hack to avoid overfitting. It is a frequency filter: it keeps the frequencies you have learned so far and discards the ones you have not yet reached.

What the theory says

The cleanest explanation comes from the neural tangent kernel (NTK) view of training. In the wide-network, small-learning-rate regime, training a network by gradient descent is approximately equivalent to kernel regression with a fixed kernel, the NTK. The dynamics of the residual error then decompose along the eigenfunctions of that kernel.

For a target function $f^{*}$ and a network prediction $f_{t}$ at training time $t$ , the error along each eigenfunction decays at a rate set by the corresponding eigenvalue. Concretely, if $λ_{k}$ is the eigenvalue for the $k$ -th eigenfunction, the component of the error in that direction shrinks roughly like:

e_{k} (t) \approx e_{k} (0) exp (- λ_{k} t)

The key fact is that for standard architectures, $λ_{k}$ decays as frequency increases. Low-frequency eigenfunctions have large eigenvalues and therefore fast decay. High-frequency eigenfunctions have small eigenvalues and therefore slow decay. The network is not choosing to ignore high frequencies. Gradient descent simply moves much faster in the low-frequency directions.

A Fourier view makes the same point in a more familiar language. Write a one-dimensional target as a sum of sinusoids:

f^{*} (x) = k \sum a_{k} sin (k x + ϕ_{k})

Spectral bias says that the coefficients $a_{k}$ for small $k$ are recovered early in training, while the coefficients for large $k$ are recovered late, if at all. The rate at which each coefficient is learned falls off with $k$ , and the falloff is steep for common activations such as tanh and ReLU in their standard parameterisations.

This is why the NTK picture is so useful even if you never compute a kernel in practice. It turns a vague observation, "networks learn smooth things first", into a quantitative claim about learning rates per frequency, and it predicts the ordering you will see in experiments.

A concrete experiment you can run

The effect is easy to see with a tiny regression problem. Take a target that mixes a low frequency and a high frequency, for example:

f^{*} (x) = sin (x) + 0.5 sin (20 x)

Train a small multilayer perceptron on samples of this function and plot the prediction at several checkpoints. Early on, the prediction looks almost exactly like $sin (x)$ . The $sin (20 x)$ term is invisible. Much later, the small rapid oscillations start to appear, and only after long training does the fit become tight.

Here is a minimal PyTorch sketch of that experiment:

import torch
import torch.nn as nn

torch.manual_seed(0)
x = torch.linspace(-3.14, 3.14, 512).unsqueeze(1)
y = torch.sin(x) + 0.5 * torch.sin(20 * x)

model = nn.Sequential(
    nn.Linear(1, 128), nn.Tanh(),
    nn.Linear(128, 128), nn.Tanh(),
    nn.Linear(128, 1),
)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()

for step in range(20000):
    opt.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    opt.step()
    if step in (100, 1000, 5000, 19999):
        print(step, float(loss))

If you snapshot the prediction at those steps, you will see the story directly: the smooth component first, the rapid component much later. Try changing the high frequency from 20 to 5 and the gap shrinks. Try changing it to 50 and the gap grows dramatically. Frequency, not just amplitude, controls the learning speed.

Why this helps generalisation, and why it hurts

Spectral bias is a double-edged sword.

On the helpful side, most real-world signals are dominated by low frequencies. Natural images, speech, sensor traces, and tabular relationships usually have most of their energy in smooth, low-frequency structure, with high-frequency energy coming disproportionately from noise. A learner that fits low frequencies first therefore gets a useful inductive bias for free: it finds the signal before it starts memorising the noise. This is one reason heavily overparameterised networks can generalise despite having enough capacity to fit pure noise. They could fit the noise, but gradient descent reaches the smooth solution first, and early stopping or weight decay can keep it there.

On the harmful side, spectral bias is a real limitation whenever the signal itself is high frequency. Three cases come up repeatedly in practice:

Sharp boundaries and edges. A classifier whose decision boundary has fine structure will learn a smoothed version of that boundary first, and may never resolve the detail within a practical training budget.
High-frequency regression targets. Problems like learning a rapidly oscillating physical field, a fine texture, or a signal with sharp peaks are intrinsically hard for a plain MLP, no matter how many parameters you add.
Coordinate-based representations. Neural radiance fields and implicit representations map coordinates to colours or densities. A plain MLP applied to raw coordinates struggles with fine detail for exactly this reason, which is why positional encoding became standard in that literature.

In each case, the failure looks like underfitting, but adding capacity or training longer only helps slowly. The bottleneck is the learning rate per frequency, not the number of parameters.

How practitioners fight spectral bias

Once you see the problem as frequency-dependent learning speed, the standard fixes make sense as ways to rebalance the spectrum.

Fourier features and positional encodings. Instead of feeding raw coordinates $x$ into the network, feed a set of sinusoids at multiple frequencies:

γ (x) = [sin (2^{0} π x), cos (2^{0} π x), \dots, sin (2^{L} π x), cos (2^{L} π x)]

This lifts high-frequency variation into the input itself, so the network no longer has to build rapid oscillations out of smooth activations. The network can then fit high-frequency targets far faster. The cost is that you also make it easier to fit high-frequency noise, so the number of frequency bands $L$ becomes a tuning knob that trades detail against overfitting.

Activation and initialisation choices. The spectral decay of the NTK depends on the activation function and on how weights are initialised. Sinusoidal activations, as in SIREN networks, change the kernel spectrum directly and are designed for exactly this kind of high-frequency fitting. More generally, scaling the input or the first-layer weights changes which frequencies the network sees early, which is why input scaling is a more consequential hyperparameter than it first appears.

Multi-scale and residual structure. Architectures that process the input at several scales, or that add high-frequency detail on top of a smooth base prediction, effectively give the network separate pathways for different frequency bands. Each pathway can then learn at its own speed rather than waiting for a single pathway to work through the spectrum in order.

Training-time frequency control. Curriculum strategies that start with smoothed targets and gradually add detail, and regularisers that penalise high-frequency components of the prediction, both exploit the same ordering. They are not fighting spectral bias so much as using it deliberately: decide which frequencies matter, and let the training dynamics do the filtering.

Limitations and common pitfalls

Spectral bias is a powerful lens, but it is easy to over-apply. Keep four cautions in mind.

First, the NTK analysis is exact only in an idealised regime: very wide networks, small learning rates, and parameters that barely move from initialisation. Real networks leave that regime, and feature learning changes the kernel during training. The frequency ordering usually survives in practice, but the precise decay rates do not transfer literally.

Second, frequency is defined relative to the input geometry. On images, graphs, or text, "frequency" means something different from the one-dimensional Fourier picture, and the right basis is the eigenbasis of the relevant kernel or graph Laplacian, not a plain sinusoid. The intuition carries over, but the maths needs the right basis for the domain.

Third, spectral bias does not mean high frequencies are impossible to learn. It means they are slower. Given enough training, capacity, and the right input representation, networks do fit high-frequency targets. A claim that a network "cannot" learn something because of spectral bias is almost always too strong. The correct claim is about cost: how much longer, and with what representation, it takes.

Fourth, do not confuse spectral bias with a preference for simple functions in general. A function can be simple in description length yet high frequency, and a function can be low frequency yet complicated. Spectral bias is specifically about the frequency spectrum under the training dynamics, not about Kolmogorov complexity or parameter count.

The takeaway

Neural networks do not learn functions uniformly. They learn them frequency by frequency, smooth structure first and fine detail later, because gradient descent moves fastest in the large-eigenvalue, low-frequency directions of the training kernel. That ordering explains a surprising amount: why early stopping works, why overparameterised models can still generalise, and why plain MLPs struggle with sharp detail unless you give them Fourier features or a multi-scale structure.

The practical rule is simple. If your target is smooth, spectral bias is working for you, and you should protect it with early stopping and modest capacity. If your target has essential high-frequency structure, you need to rebalance the spectrum explicitly, with input encodings, activations, or architecture, rather than just training longer and hoping the wiggles arrive in time.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial