Regularisation and normalisation

Your training loss falls. Your validation loss falls at first, then rises. Adding a larger model makes the training result even better and the validation result worse.

That pattern is often called overfitting: the fitted model performs much better on the examples it learned from than on representative held-out examples. The response is not to switch on every available technique. First protect the evidence. Then choose one intervention with a reason.

Why this matters

Real training data contains limited coverage, noise, accidental correlations, and repeated patterns. A flexible network can reduce training loss by relying on details that will not remain useful after deployment.

Regularisation gives the learning process a deliberate preference: smaller weights, incomplete activation paths, or task-valid input variation. Normalisation keeps internal numbers on a controlled coordinate system. Both can affect optimisation and validation results, but they do different jobs.

Plain-English mental model

Step 1 — Diagnose before adding a constraint

Overfitting is an evidence pattern, not a feeling. Look for a persistent gap between the training measure and a representative validation measure. Then check easier explanations first:

1. Is the split independent, representative, and free from duplicated entities or future information?
2. Does validation use the same deterministic preprocessing contract intended for deployment?
3. Did evaluation call model.eval() and disable gradient tracking separately?
4. Is the chosen measure hiding an important error slice or class?
5. Does a simpler baseline show the same gap?

If the data or evaluation path is wrong, more regularisation can hide the symptom without fixing the experiment.

A useful comparison changes one important factor from a reproducible baseline. Keep the data split, checkpoint rule, optimiser schedule, and evaluation measures recorded. Select the factor using validation data, not the final test set.

Step 2 — Three regularisation levers

These methods do not mean “make the model worse.” They make some training solutions less convenient, hoping that a more transferable solution wins.

Dropout: make activation paths unreliable during training

DropoutDropoutA training-only regularisation method that randomly replaces some activations with zero and scales the survivors. At evaluation time, the dropout layer passes every activation through. randomly replaces activations with zero during training. Its setting p is the probability of dropping, not keeping, an activation.

If p = 0.5, each activation has a 50% chance of surviving. PyTorch scales each survivor by 1 / (1 - p), which is 2 here. That scaling keeps the expected activation value unchanged across many random masks.

Suppose the activations are [2, 4, 6, 8] and one fixed teaching mask is [1, 0, 0, 1]:

training = activation × mask ÷ (1 − p)
         = [4, 0, 0, 16]

evaluation = [2, 4, 6, 8]

The original mean and this particular training-mask mean both happen to be 5. The general guarantee is about the expectation over masks; one arbitrary mask need not preserve the mean exactly.

Dropout makes repeated training forward passes random. During evaluation, a PyTorch dropout layer becomes an identity operation. That change is controlled by model mode:

model.train()  # dropout masks activations
model.eval()   # dropout passes activations through

model.eval() does not disable gradient tracking. torch.no_grad() does not switch dropout or BatchNorm to evaluation behaviour. A correct evaluation path normally needs both:

model.eval()
with torch.no_grad():
    validation_prediction = model(validation_features)

The dropout probability is a hyperparameter. A larger p is a stronger disturbance, not an automatic improvement. Record where dropout is applied and compare its validation effect.

Weight decay: nudge selected parameters towards zero

Weight decayWeight decayAn optimiser rule that nudges selected parameters towards zero on each update. Its strength is a tunable regularisation setting, not a guarantee of better generalisation. shrinks selected parameters a little on each optimiser update. For a simple decoupled update, the decay part looks like:

weight after decay = weight × (1 − learning rate × decay strength)

Smaller parameters can reduce reliance on a few extreme numeric weights. Zero decay means no shrinkage. Too much decay can prevent the model fitting a real pattern.

AdamW keeps the decay step separate from the gradient moments used for Adam’s adaptive update. That is different from placing a squared-weight penalty inside the loss and passing its gradient through Adam’s adaptive scaling. The toy calculation below shows only where the decay term enters; it is not a full implementation of either optimiser.

weight = 2.0
data_gradient = 0.5
adaptive_scale = 0.2
learning_rate = 0.1
decay = 0.01

coupled = weight - learning_rate * adaptive_scale * (
    data_gradient + decay * weight
)
decoupled = (
    weight * (1 - learning_rate * decay)
    - learning_rate * adaptive_scale * data_gradient
)

print(f'coupled toy update: {coupled:.4f}')
print(f'decoupled toy update: {decoupled:.4f}')

Decide which parameters receive decay. A common starting policy decays matrix-shaped weights while excluding one-dimensional bias and normalisation scale/shift parameters. That is a policy to test, not a law. Make the parameter groups explicit so a reviewer can see the choice.

Data augmentation: vary inputs without changing the answer

Data augmentationData augmentationCreating varied training inputs with transformations that should preserve the target for the intended task, while keeping validation and test preprocessing deterministic. creates varied training inputs with transformations that should preserve the target for the intended task.

That last phrase is the contract. A horizontal flip may preserve “contains a cat.” It may reverse a left-versus-right label, make text unreadable, or move an object box unless the box is transformed too.

InvarianceInvarianceThe desired property that a model's relevant output stays stable when an input changes in a way that should not change the answer for the task. means the relevant output should stay stable under a change that should not alter the task answer. Augmentation teaches a chosen invariance by showing transformed training examples with consistent targets.

This small pure-Python example flips a three-by-three mark. For this deliberately narrow label—“contains a diagonal mark”—the label remains valid:

image = ['#..', '.#.', '..#']
label = 'contains a diagonal mark'

flipped = [row[::-1] for row in image]

print('\n'.join(flipped))
print('label:', label)

Random augmentation belongs in the training path. Validation and test preprocessing should be deterministic and should match the intended deployment contract. If an input has a paired target—such as a segmentation mask, bounding box, keypoint, or audio timestamp—transform the pair consistently.

An invalid transform is structured label noise. Audit transformed examples visually and by target type before trusting a validation score.

Step 3 — Normalisation starts with an axis

Normalisation changes a collection of values using a statistic such as a mean, variance, or root mean square. Before naming a layer, ask two questions:

1. Which values share the statistic—examples in a batch, or features inside one example?
2. Does evaluation use current-input statistics or stored training statistics?

In the formulas below, epsilon is a tiny positive value that prevents division by zero. Learnable scale and shift parameters let the layer restore or reshape useful magnitudes after normalisation.

BatchNorm: across examples for each feature

Batch normalisationBatch normalisationA layer that normalises each feature using statistics calculated across training examples in a mini-batch, then applies learned scale and shift parameters. It normally uses stored running statistics during evaluation., usually called BatchNorm, treats each feature separately. During training it calculates that feature’s mean and variance across examples in the mini-batch:

normalised value = (value − batch feature mean)
                   ÷ sqrt(batch feature variance + epsilon)

output = learned scale × normalised value + learned shift

The selected example therefore depends on its batch neighbours during training. Replace one neighbour with a large outlier and the selected example’s output can change.

By default, a PyTorch BatchNorm module also updates running estimates during training and uses those estimates during evaluation. This makes model.train() and model.eval() semantically important. Tiny or non-representative batches can produce noisy training statistics. Accumulating gradients across micro-batches does not make BatchNorm see one larger activation batch; each forward pass still has its own batch statistics.

For an image tensor, BatchNorm commonly has one learned scale and shift per channel. The exact axes depend on the BatchNorm module and input shape, so read the module’s shape contract.

LayerNorm: centre and scale within one example

Layer normalisationLayer normalisationA layer that centres and scales an example using statistics calculated across its selected feature dimensions, followed by learned element-wise scale and shift parameters., or LayerNorm, calculates its mean and variance over selected feature dimensions inside each example:

normalised value = (value − example mean)
                   ÷ sqrt(example variance + epsilon)

output = learned element-wise scale × normalised value
         + learned element-wise shift

Changing another example in the mini-batch does not change this example’s LayerNorm statistics. PyTorch LayerNorm uses current-input statistics in both training and evaluation. It still has learned parameters, but it does not need BatchNorm-style running estimates.

If every feature in the normalised group is identical, the centred values are zero. Epsilon keeps the division defined, but the layer cannot invent variation that was absent.

RMSNorm: scale without subtracting the mean

Root mean square normalisationRMS normalisationA layer that divides an example by the root mean square of values across selected feature dimensions, then applies a learned scale. Unlike layer normalisation, it does not subtract the mean., usually called RMSNorm, calculates the root mean square across selected feature dimensions:

root mean square = sqrt(mean(value²) + epsilon)
output = learned scale × value ÷ root mean square

RMSNorm does not subtract the mean and does not add the LayerNorm-style learned bias in PyTorch’s module. A uniform shift of every feature therefore still changes the normalised direction and values.

Like LayerNorm, its statistics come from the current example rather than other mini-batch examples. Choose it because that behaviour fits the architecture and tested baseline, not because it is newer or has a shorter formula.

Worked example — one row, three normalisations

Take a selected row [1, 2, 4, 8] inside a batch containing the same row shifted upward by one, two, and three. Ignore learnable scale/shift and the tiny epsilon for this teaching calculation.

For LayerNorm, the selected row mean is 3.75 and its population standard deviation is about 2.681. Centre and divide each value to get approximately:

LayerNorm: [-1.026, -0.653, +0.093, +1.585]

For RMSNorm, the selected row’s root mean square is about 4.610. Divide without centring:

RMSNorm: [+0.217, +0.434, +0.868, +1.735]

For BatchNorm, each feature column is an evenly spaced sequence of four values. The selected row is the lowest value in every column, so every feature becomes approximately -1.342:

BatchNorm selected row: [-1.342, -1.342, -1.342, -1.342]

from math import sqrt

selected = [1.0, 2.0, 4.0, 8.0]
batch = [[value + shift for value in selected] for shift in range(4)]

def mean(values):
    return sum(values) / len(values)

def standardise(values):
    centre = mean(values)
    spread = sqrt(mean([(value - centre) ** 2 for value in values]))
    return [(value - centre) / spread for value in values]

layer = standardise(selected)
rms = sqrt(mean([value ** 2 for value in selected]))
rms_output = [value / rms for value in selected]

batch_output = []
for feature in range(len(selected)):
    column = [row[feature] for row in batch]
    batch_output.append(standardise(column)[0])

print('batch:', [f'{value:+.3f}' for value in batch_output])
print('layer:', [f'{value:+.3f}' for value in layer])
print('rms:  ', [f'{value:+.3f}' for value in rms_output])

Now use the interactive version. Keep Four nearby examples selected and change only the uniform shift. LayerNorm output stays fixed because centring removes that shift. RMSNorm changes because it does not centre. Then select the outlier batch: only the BatchNorm result for the selected row changes.

The lab fixes learned scale and shift to identity so you can see the statistics. A production layer learns those parameters and includes epsilon in the denominator.

Step 4 — Assemble an explicit PyTorch policy

This small classifier places LayerNorm after the first linear projection, then applies an activation and dropout. The model uses no BatchNorm, so its activation statistics do not depend on other examples in the mini-batch.

The optimiser groups matrix-shaped weights separately from one-dimensional biases and normalisation parameters. This is one visible starting policy; a real architecture may require named exceptions.

import torch
from torch import nn

class RegularisedMLP(nn.Module):
    def __init__(self, input_features: int, hidden_features: int, classes: int):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_features, hidden_features),
            nn.LayerNorm(hidden_features),
            nn.GELU(),
            nn.Dropout(p=0.2),
            nn.Linear(hidden_features, classes),
        )

    def forward(self, features):
        return self.network(features)

model = RegularisedMLP(input_features=8, hidden_features=16, classes=3)

decay_parameters = []
no_decay_parameters = []
for parameter in model.parameters():
    group = decay_parameters if parameter.ndim >= 2 else no_decay_parameters
    group.append(parameter)

optimizer = torch.optim.AdamW(
    [
        {'params': decay_parameters, 'weight_decay': 0.01},
        {'params': no_decay_parameters, 'weight_decay': 0.0},
    ],
    lr=3e-4,
)

print('decayed tensors:', len(decay_parameters))
print('non-decayed tensors:', len(no_decay_parameters))

For this exact model, the expected final two lines are:

There are two linear weight matrices. The two linear biases plus LayerNorm’s learned scale and bias are one-dimensional and enter the no-decay group.

Test mode behaviour directly:

features = torch.ones(1, 8)

model.train()
training_a = model(features)
training_b = model(features)  # usually differs because dropout resamples

model.eval()
with torch.no_grad():
    evaluation_a = model(features)
    evaluation_b = model(features)

assert torch.equal(evaluation_a, evaluation_b)

“Usually differs” is deliberately not an assertion: two random dropout masks can occasionally match. The evaluation equality is the deterministic property this example tests on a fixed device with deterministic model operations.

Step 5 — Choose with a controlled ablation

An ablation is a comparison that removes or changes one component to measure what evidence changes. It does not prove a universal cause, but it is more informative than changing the whole recipe.

Start with these questions:

• Dropout: Which activations receive it? What is p? Is the validation path definitely in evaluation mode?
• Weight decay: Which parameters are included? What is the decay strength? Is the optimiser AdamW or a different rule?
• Augmentation: Which task invariance does each transform represent? Are paired targets transformed consistently?
• Normalisation: Which axes share statistics? Are those statistics current-input or running? What happens with the deployment batch shape?

Compare training curves, the declared validation measure, relevant error slices, and operational measures such as latency. A smaller training–validation gap caused by much worse performance on both sets is not automatically a success.

Real-world AI engineering example

Common mistakes

Knowledge check

Practice activity

Exercises

Key takeaways

Flashcards

One-page sketchnote summary

train: augment → model.train()
validate: deterministic → model.eval()
compare one declared change

Resources