Overfitting & regularisation

A model that scores 99% on training data but 70% on new data has memorised rather than learned. This is overfitting, and with millions of parameters, deep networks are especially prone to it. Regularisation is the toolbox of techniques that push a model towards patterns that generalise.

Diagnosing the problem

Always keep three splits: train (fit weights), validation (tune choices) and test (final, one-time score). Plot train and validation loss per epoch:

• Both falling together – keep training.
• Train falling, validation rising – overfitting has begun.
• Both high – underfitting; the model is too weak or undertrained.

The gap between train and validation accuracy is your overfitting meter.

The regularisation toolbox

Technique	How it helps	PyTorch
More / augmented data	More variety, harder to memorise	torchvision transforms
Dropout	Randomly zeros activations during training so no neuron can be relied on alone	nn.Dropout(p=0.3)
Weight decay	Penalises large weights, favouring simpler functions	AdamW(..., weight_decay=1e-2)
Early stopping	Stops when validation stops improving	a few lines of code
Batch normalisation	Normalises layer inputs per batch; stabilises training and adds mild noise	nn.BatchNorm1d/2d
Smaller model	Fewer parameters to memorise with	fewer layers or units
Label smoothing	Prevents over-confident predictions	CrossEntropyLoss(label_smoothing=0.1)

Dropout and batch norm in a model

import torch.nn as nn

class RegularisedNet(nn.Module):
    def __init__(self, n_in: int, n_classes: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_in, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, n_classes),
        )

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

Dropout is only active in model.train() mode; in model.eval() it passes everything through (PyTorch scales activations during training so the expected values match). Batch norm uses batch statistics while training and stored running averages during evaluation. This is why forgetting model.eval() makes predictions noisy and inconsistent.

A practical note: batch norm needs more than one example per batch in training mode, and very small batches (under about 8) give poor statistics. For such cases, or for transformers, LayerNorm is used instead.

Weight decay

Weight decay adds a pull towards zero on every weight at each step. With AdamW it is "decoupled" from the adaptive update, which works better than adding an L2 term to the loss when using Adam.

import torch

opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.05)

Typical values range from 1e-4 to 0.1; tune on validation data.

Early stopping with checkpointing

import copy

best_loss, best_state, patience, bad_epochs = float("inf"), None, 5, 0

for epoch in range(100):
    train_one_epoch(model, train_dl)
    val_loss = validate(model, val_dl)
    if val_loss < best_loss - 1e-4:
        best_loss, bad_epochs = val_loss, 0
        best_state = copy.deepcopy(model.state_dict())
    else:
        bad_epochs += 1
        if bad_epochs >= patience:
            print(f"Stopping at epoch {epoch + 1}")
            break

model.load_state_dict(best_state)   # roll back to the best epoch

Here train_one_epoch and validate are your own helpers from the previous lesson. The important part is restoring the best weights, not the last ones.

A worked scenario

A Surat textile exporter trains a classifier on 3,000 fabric photos across 6 weave types. After 30 epochs train accuracy is 99% but validation is 78%. A sensible fix order:

1. Add augmentation (flips, small rotations, colour jitter) – biggest win for images.
2. Switch to a pretrained backbone (transfer learning) instead of training from scratch.
3. Add dropout 0.3 before the classifier and weight decay 0.01.
4. Use early stopping with patience 5.

This typically closes most of the gap. Collecting more labelled photos of the confused classes is often the best long-term fix.

Common mistakes

• Tuning hyperparameters on the test set, which leaks information and inflates results.
• Applying augmentation to validation or test data.
• Stacking every technique at maximum strength, causing underfitting.
• Using dropout directly after the final layer that produces logits.

Key takeaways

• Watch the train–validation gap; it tells you whether to regularise or add capacity.
• Data (real or augmented) is the most powerful regulariser.
• Dropout, weight decay, batch norm and early stopping are cheap and composable.
• Always switch to model.eval() for validation and inference.