Search Tech Journey

Find topics, journeys and posts

6-month learning plan101 / 130
back to blog
mladvanced 15m read

S101 · Regularization in DL — Dropout, BatchNorm, Weight Decay

Keeping deep nets from memorising.

Module M12: Deep Learning · Session 101 of 130 · Track: ML 🧊

What you'll be able to do after this session

  • Explain Regularization in DL to a friend in one minute, out loud, no notes.
  • Recognise it when you see it in real code / systems / papers.
  • Complete the hands-on exercise at the end without looking up help.

Prerequisites

If you can't answer these in 30 seconds each, redo the linked session first:


Pre-read (skim before the session)


(a) Intuition · 5 min

The one-paragraph mental model. Why this concept exists, what problem it solves, and the analogy that makes it stick.

In one sentence: Keeping deep nets from memorising.

Imagine you're studying for an exam by memorising last year's worked solutions verbatim. If the professor changes a single variable name, you're toast. That's what an over-parameterised neural network does when you let it train too long on too little data — it memorises the training set instead of learning the underlying pattern. Regularization is the collection of tricks we use to force the network to generalise instead of memorise.

There are three tricks you'll see everywhere. Weight decay (a.k.a. L2 regularization) whispers to the optimizer, "prefer small weights" — because huge weights are how a net encodes memorised outliers. Dropout randomly switches off some neurons on every forward pass during training, so no single neuron becomes a load-bearing wall; the network has to build redundant, robust representations. BatchNorm normalises each layer's activations to zero mean and unit variance, which both stabilises training and acts as a mild regularizer by injecting a bit of noise from the mini-batch statistics.

Gotcha to carry forever: Dropout and BatchNorm behave differently in .train() vs .eval() mode. Forgetting to call model.eval() before validation is the single most common bug in deep-learning tutorials — you'll get wildly inconsistent accuracy numbers and spend an hour blaming your data loader.


(b) Visual walkthrough · 15 min

Diagrams, worked examples, step-by-step visuals.

Here's the mental model of what each regularizer does to a network during a single forward pass:

Worked example 1 — Dropout arithmetic. Suppose a hidden layer outputs [2.0, 4.0, 6.0, 8.0] and dropout probability p = 0.5. In training mode PyTorch does two things: (1) it randomly zeroes each element with probability 0.5, say the mask is [1, 0, 1, 1], giving [2, 0, 6, 8]; (2) it scales the survivors by 1/(1−p) = 2, giving [4, 0, 12, 16]. This "inverted dropout" scaling means the expected value of each activation stays the same, so you don't need any special logic at inference time — you just turn dropout off.

Worked example 2 — Weight decay. Your loss is normally L = CrossEntropy(y_pred, y_true). With weight decay λ = 1e-4, the actual loss becomes L_total = L + 1e-4 · Σ w² where the sum runs over every weight in the model. The gradient of the penalty w.r.t. w is 2 · 1e-4 · w, so at every step SGD subtracts a tiny bit proportional to the weight itself — the weights are literally "decayed" toward zero every step, unless the data-fitting gradient pulls them the other way.

TechniqueTrainingInferenceTypical hyperparam
Weight decayAdds λ·‖w‖² to lossNothing (weights already small)λ = 1e-4 to 1e-2
DropoutRandomly zeros neurons + scalesIdentity (off)p = 0.1 to 0.5
BatchNormUses batch mean/varUses running mean/var learned in trainingmomentum = 0.1

Notice the pattern: every regularizer has a training-time behaviour and a different inference-time behaviour. That's why the .train() / .eval() toggle exists.


(c) Hands-on · 20 min

Code you type and run. Not pseudo-code — real, runnable, copy-pasteable.

# regularization_demo.py — MNIST with vs without regularization
import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
 
device = "cuda" if torch.cuda.is_available() else "cpu"
train = DataLoader(datasets.MNIST("./data", train=True, download=True,
                                  transform=transforms.ToTensor()),
                   batch_size=128, shuffle=True)
test  = DataLoader(datasets.MNIST("./data", train=False,
                                  transform=transforms.ToTensor()),
                   batch_size=256)
 
class Net(nn.Module):
    def __init__(self, use_reg=True):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 512)
        self.bn1 = nn.BatchNorm1d(512) if use_reg else nn.Identity()
        self.do  = nn.Dropout(0.5)     if use_reg else nn.Identity()
        self.fc2 = nn.Linear(512, 10)
    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.do(F.relu(self.bn1(self.fc1(x))))
        return self.fc2(x)
 
def run(use_reg, wd):
    model = Net(use_reg).to(device)
    opt = torch.optim.SGD(model.parameters(), lr=0.05,
                          momentum=0.9, weight_decay=wd)
    for epoch in range(3):
        model.train()                              # critical
        for x, y in train:
            x, y = x.to(device), y.to(device)
            opt.zero_grad(); F.cross_entropy(model(x), y).backward(); opt.step()
        model.eval()                               # critical
        correct = 0
        with torch.no_grad():
            for x, y in test:
                x, y = x.to(device), y.to(device)
                correct += (model(x).argmax(1) == y).sum().item()
        print(f"reg={use_reg} wd={wd} epoch={epoch} acc={correct/len(test.dataset):.4f}")
 
print("--- NO regularization ---");  run(False, 0.0)
print("--- WITH regularization ---"); run(True,  1e-4)

Checklist while it runs:

  1. Watch that the unregularized run's training loss drops faster but test accuracy stalls or drifts.
  2. Confirm the regularized run has slightly slower training loss but higher final test accuracy.
  3. Comment out model.eval() and rerun — you'll see test accuracy jitter wildly. That's the bug.
  4. Print model.bn1.running_mean[:5] after training — those are the values used at inference.
  5. Try batch_size=1 and watch BatchNorm blow up — BN needs a batch.

Try this modification: Set nn.Dropout(0.9) (drop 90% of neurons). Training will barely learn — you've regularized so hard the network has no capacity left. Now dial it down and find the sweet spot for MNIST.


(d) Production reality · 10 min

How this shows up in real systems, gotchas, war stories, common bugs.

In production, regularization is where "it worked on my laptop, it broke in staging" comes from. A few war stories.

The frozen-BN-in-fine-tuning bug. When you fine-tune a vision model (say ResNet-50 pretrained on ImageNet) on your own small dataset, the BatchNorm layers' running statistics were learned on ImageNet. If you leave them in .train() mode, tiny mini-batches from your dataset will corrupt those statistics and destroy the pretrained features. Every serious CV team either freezes BN or replaces it with GroupNorm. Detectron2, MMDetection, and torchvision all ship helpers for exactly this.

The dropout-at-inference bug. People deploy PyTorch models via TorchScript or ONNX and forget to call model.eval() before tracing. The exported graph then contains random dropout at inference, meaning the same input gives different predictions every call — a nightmare to debug because it looks like a "flaky model". Always trace/export after .eval().

Weight decay ≠ L2 for Adam. This is the subtle one. Vanilla L2 (adding λ‖w‖² to the loss) interacts badly with Adam because Adam's adaptive learning rates make the effective decay uneven across parameters. Loshchilov & Hutter's AdamW (2019) decouples the decay from the gradient, and it's now the default for transformer training. If you're using plain torch.optim.Adam(..., weight_decay=1e-2) for a large model, you're probably doing it wrong — switch to torch.optim.AdamW.

At scale. Meta's FSDP blog notes that when training multi-billion-parameter models, they combine weight decay with gradient clipping and mixed-precision training — each individually mild, but together they stabilise runs that would otherwise diverge in the first thousand steps. Regularization in the LLM era is less about "generalisation on a small dataset" and more about "keep the loss curve smooth on a giant one".


(e) Quiz + exercise · 10 min

  1. In one sentence, what does dropout do during training, and what does it do at inference time?
  2. Why must you call model.eval() before validation when the model contains BatchNorm?
  3. Given p = 0.2 dropout, by what factor does PyTorch scale the surviving activations?
  4. What is the difference between L2 regularization and AdamW-style weight decay, and why does it matter for transformers?
  5. Name two production symptoms that indicate you forgot model.eval() before exporting.

Stretch (connects to S099 — Optimizers): SGD with momentum + weight decay is mathematically equivalent to what other regularization scheme, and how does that equivalence break down for Adam?


Explain-out-loud test

If you can't teach these 3 things to a friend without notes, redo the session:

  1. What is Regularization in DL? (one sentence, no jargon)
  2. When would you use it? (one concrete example)
  3. When would you NOT use it? (one anti-pattern)

What comes next


Part of a 130-session evergreen learning series. Session structure: (a) intuition · (b) visual walkthrough · (c) hands-on · (d) production reality · (e) quiz + exercise. Duration: 60 minutes.

More from M12 · Deep Learning

all modules →
  1. 096Perceptron & Activation Functions
  2. 097Multi-Layer Perceptron — Forward Pass
  3. 098Backpropagation — Derived by Hand on a 2-Layer Net
  4. 099Optimizers — SGD, Momentum, Adam, RMSprop
  5. 100PyTorch Fundamentals — Tensors, Autograd, nn.Module
  6. 102CNNs — Convolution, Pooling, ImageNet Architectures