Regularization in DL — Dropout, BatchNorm, Weight Decay

Why this session exists

Training fit and predictive generalization are different goals. Dropout changes training-time activations, normalization changes feature statistics and parameterization, and weight decay changes the update. They can help, hurt or have little effect depending on data and architecture. A wide network does not necessarily memorize every possible dataset, and successful deep learning does not require all three techniques in every block.

You will derive inverted dropout's expectation and variance, distinguish normalization axes and train/eval state, derive coupled versus decoupled decay, and run a controlled CPU ablation with validation-only selection. The final test is not a knob-tuning surface. A smaller train/validation gap is not automatically a better model.

Prerequisites

Use S099 optimizer equations, S100 module modes and buffers and S097 matrix shapes. Keep data loss separate from any explicit penalty and distinguish online training loss from evaluation-mode loss on training rows.

(a) Intuition · 5 min

A study group that cannot always rely on the same participant must distribute useful knowledge. Dropout similarly disrupts specific co-adaptations by randomly removing activation contributions during training. The analogy does not imply that a dropped unit learns in that pass or that removing more participants always improves the group.

Weight decay directly shrinks selected parameters; it can discourage some large-norm solutions but does not universally order function complexity, especially with scale-invariant normalization. Normalization rescales intermediate coordinates and can help optimization; BatchNorm's batch-dependent statistics also introduce noise. LayerNorm has no cross-example batch noise, so describing every normalizer as the same regularizer is inaccurate.

When you need which

Compare training and validation losses under the same evaluation mode, data transform and reduction before diagnosing a gap. Low training loss and worse validation loss can reflect overfitting, distribution shift, class composition, leakage or label quality. If both are poor, inspect optimization and targets as well as capacity/regularization. A flat curve from step one is not a command to add BatchNorm.

Choose one intervention at a time where possible: label-preserving augmentation, a smaller architecture, decay, dropout or an early-stopping rule. Use validation for selection and assess the frozen choice on test. Fixed epochs are also legitimate when specified beforehand; early stopping is not universally mandatory. More data helps only to the extent its labels and sampling support the target problem; it does not add no bias under all collection procedures.

The story of when we learned to regularise

The dropout paper explains sampling thinned networks and approximating their prediction average with an unthinned network. Layer normalization's primary paper proposes per-example normalization suitable for recurrent computation. AdamW's primary paper separates weight decay from adaptive gradient normalization. These contributions are distinct; they do not support claims that every transformer adopted AdamW in 2019 or that one technique universally “enabled” a particular architecture.

The three regularisers in a picture

The decay arrow is not an extra penalty added to the loss for AdamW. The diagram shows one possible block order, not a universal canonical order. Pre-activation residual designs and other normalizer placements have different contracts.

Dropout — the training-time / inference-time asymmetry

Let drop probability be $p<1$, keep probability $q=1-p$ and mask $m\sim\operatorname{Bernoulli}(q)$. Inverted dropout returns $\tilde a=ma/q$ during training and $a$ during ordinary evaluation. For fixed activation $a$,

Derive the mean from $E[m]=q$, the second moment from $m^2=m$, then subtract the squared mean. For $a=2,p=0.25$, the output is 0 with probability 0.25 and $8/3$ with probability 0.75; the mean is 2 and variance $4/3$. As $p$ approaches one, the surviving magnitudes and variance grow. PyTorch permits p=1 as an all-zero training special case; the expectation-preservation derivation does not apply there.

Expectation preservation applies to that activation and a following affine sum. It does not mean the complete nonlinear network in eval mode equals the exact expected prediction of all masked networks: generally $E[f(\tilde a)]\ne f(E[\tilde a])$. “An ensemble for free” is an intuition about shared-weight subnetworks, not independently trained models or exact averaging. Ordinary nn.Dropout drops elements; spatial/channel variants have different mask structure.

In ordinary eval mode dropout is identity, not “all neurons become positive.” Existing zero activations remain zero. Deliberate Monte Carlo dropout is another inference policy; it can keep selected dropout active while avoiding updates to BatchNorm statistics. Stochastic inference is not inherently invalid, but it must be intentional and evaluated.

BatchNorm vs LayerNorm vs GroupNorm

For values $x_i$ in a normalization group of size $m$,

This standardizes selected marginals; it is not full whitening because feature covariances are not removed. Epsilon and learned affine parameters also mean outputs are not generally exactly zero-mean/unit-variance.

LayerNorm can normalize several specified trailing dimensions, not always a whole example. GroupNorm requires channels divisible by the group count and usually has per-channel affine parameters. LayerNorm has affine parameters per normalized element; even when reduction groups match, affine parameterization can differ.

PyTorch BatchNorm details: train forward uses the biased variance (unbiased=False), while the running variance update uses the unbiased estimate. With momentum $m_b$, running state updates as $(1-m_b)old+m_b\,batch$, the opposite coefficient naming intuition from optimizer momentum. track_running_stats=False uses batch statistics even in eval. One value per channel in training is insufficient for the standard BatchNorm path, whereas one image with multiple spatial values can still provide several observations per channel.

RMSNorm:

It does not subtract the mean and PyTorch's module has a learned scale but no additive bias. For $x=(1,3)$, LayerNorm without affine is approximately $(-1,1)$, while RMSNorm gives approximately $(1/\sqrt5,3/\sqrt5)$. They are not the same function with fewer parameters. LayerNorm/RMSNorm's per-token axis is convenient in sequence architectures; BatchNorm is common in CNNs; GroupNorm is useful when batch-independent channel-group statistics are desired. These are design patterns, not universal architecture prohibitions or equal-accuracy claims.

Dropout and BatchNorm interactions

Dropout changes second moments even when it preserves a conditional mean. If random activation $a$ has mean $\mu$ and variance $\sigma^2$, independent inverted dropout gives

A downstream BatchNorm can therefore see a different input variance in training and eval. Putting dropout after one BN avoids perturbing that BN's current input, but a later block's BN may still receive changed statistics. There is no theorem that using both in one block always hurts or that one ordering eliminates all mismatch. The original combined recipe is retained as a tested variant, not promoted to a universal default.

Weight decay — where it lives depends on the optimiser

For the objective $L_{data}+\lambda\|w\|^2/2$, plain SGD gives

Thus additive L2 and multiplicative shrinkage coincide for that plain-SGD convention. With momentum, a penalty entering the velocity is not generally equivalent to a separate shrink with the same coefficient. Coupled Adam computes moments from $g+\lambda w$; AdamW computes moments from $g$ and applies a separate $(1-\eta\lambda)w$ factor. The effect is more than simply dividing a fixed decay term by an already-fixed second moment: the penalty also changes the moment estimates.

Coupled L2-Adam is a valid optimizer for an L2-penalized objective; AdamW is a different regularization procedure. In PyTorch 2.9, Adam(decoupled_weight_decay=True) also implements the decoupled choice. Match the installed version rather than calling every Adam(weight_decay=...) use a bug.

With data gradient zero, a parameter with an explicit zero gradient tensor and no earlier moment state under AdamW shrinks by $(1-\eta\lambda)^T$ over $T$ steps. Parameters whose gradients are None can be skipped. Cumulative decay depends on the rate schedule and number of updates, not just weight_decay. A no-decay policy for biases/normalization parameters is a parameter-group decision to document, not a law.

Augmentation, explicit regularization or smaller models?

Augmentation encodes assumptions such as small translations preserving an image label. It does not create independent observed labels; an invalid transform can introduce bias, and even a valid invariance may not match deployment frequencies. Tabular augmentations are possible but require meaningful domain constraints.

Decay and dropout provide continuous controls but interact with optimization. Shrinking the model reduces parameter and activation costs but may remove needed capacity; neither variance nor quality must improve monotonically. Early stopping selects an iteration budget on validation. Choose interventions against the observed failure and operational budget instead of “weight decay always, early stopping always.”

(c) Hands-on · 25 min

This CPU experiment compares vanilla, dropout, BN, the original combined pattern, and isolated decay strengths. All variants share the initial linear weights and batch order; new BN affine parameters start at standard identity settings. Training curves are measured with the evolving stochastic model, while reported train/validation gaps use the same eval-mode checkpoint. The best validation checkpoint/variant is selected before the test is evaluated once.

import copy
import numpy as np
import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
 
torch.set_num_threads(1); torch.manual_seed(10)
class MLP(nn.Module):
    def __init__(self,drop=0.,bn=False):
        super().__init__()
        self.fc1=nn.Linear(64,48); self.fc2=nn.Linear(48,24); self.fc3=nn.Linear(24,10)
        self.bn1=nn.BatchNorm1d(48) if bn else nn.Identity()
        self.bn2=nn.BatchNorm1d(24) if bn else nn.Identity()
        self.drop=nn.Dropout(drop)
    def forward(self,x):
        x=self.drop(torch.relu(self.bn1(self.fc1(x))))
        x=self.drop(torch.relu(self.bn2(self.fc2(x))))
        return self.fc3(x)
 
X,y=load_digits(return_X_y=True); X=torch.tensor(X/16.,dtype=torch.float32); y=torch.tensor(y,dtype=torch.long)
dev,test=train_test_split(np.arange(len(y)),test_size=.2,stratify=y.numpy(),random_state=1)
train,val=train_test_split(dev,test_size=.25,stratify=y[dev].numpy(),random_state=2)
assert not (set(train)&set(val) or set(dev)&set(test))
base=copy.deepcopy(MLP().state_dict()); criterion=nn.CrossEntropyLoss()
def loader(ids,shuffle=False):
    return DataLoader(TensorDataset(X[ids],y[ids]),batch_size=64,shuffle=shuffle,
                      generator=torch.Generator().manual_seed(20),num_workers=0)
def epoch(model,ld,opt=None):
    model.train(opt is not None); total=0.; n=0
    with torch.set_grad_enabled(opt is not None):
        for x,yb in ld:
            if opt is not None: opt.zero_grad(set_to_none=True)
            loss=criterion(model(x),yb); assert torch.isfinite(loss)
            if opt is not None: loss.backward(); opt.step()
            total+=loss.item()*len(yb); n+=len(yb)
    return total/n
 
def weight_norm(model):
    # Norm of the concatenated matrix parameters, NOT sum of individual norms.
    return float(torch.sqrt(sum(p.detach().square().sum() for p in model.parameters() if p.ndim>1)))
 
variants=[('vanilla',0.,False,0.),('dropout',.3,False,0.),('BN',0.,True,0.),
          ('combined',.3,True,.05),('decay-.05',0.,False,.05),('decay-.5',0.,False,.5)]
results=[]
for name,drop,bn,wd in variants:
    torch.manual_seed(30); model=MLP(drop,bn)
    state=model.state_dict()
    for key,value in base.items():
        if key.startswith('fc'): state[key]=value.clone()
    model.load_state_dict(state)
    decay=[p for p in model.parameters() if p.ndim>1]
    no_decay=[p for p in model.parameters() if p.ndim<=1]
    opt=torch.optim.AdamW([{'params':decay,'weight_decay':wd},
                           {'params':no_decay,'weight_decay':0.}],lr=.005)
    train_ld=loader(train,True); best=float('inf'); curve=[]
    for ep in range(1,13):
        online=epoch(model,train_ld,opt)
        train_eval=epoch(model,loader(train)); validation=epoch(model,loader(val))
        curve.append((ep,online,train_eval,validation))
        if validation<best:
            best=validation; chosen=copy.deepcopy(model.state_dict()); chosen_ep=ep
    model.load_state_dict(chosen)
    train_eval=epoch(model,loader(train)); validation=epoch(model,loader(val))
    norm=weight_norm(model)
    results.append((validation,name,chosen,drop,bn))
    print(name,'checkpoint',chosen_ep,'eval train/val/gap',train_eval,validation,validation-train_eval,'matrix L2 norm',norm)
    print('curve epoch/online-train/eval-train/validation',curve)
_,name,state,drop,bn=min(results,key=lambda r:r[0])
winner=MLP(drop,bn); winner.load_state_dict(state)
print('validation-selected winner',name,'final test loss',epoch(winner,loader(test)))
winner.eval()
with torch.no_grad(): assert torch.equal(winner(X[:8]),winner(X[:8]))
assert all(np.isfinite(row[0]) for row in results)

Anatomy of the script

The same fixed pixel scaling applies to all partitions; it is not estimated using test rows. Initialization and loader generators are reset independently so dropout's RNG consumption does not silently change batch order. BN adds parameters and buffers, so this is a mechanism ablation, not a perfectly identical total-parameter comparison. The combined variant changes several mechanisms at once and cannot isolate which one caused its result; the separate variants help interpret it.

Online training loss mixes different parameter states and active dropout/batch statistics. It is not directly comparable with a final eval-mode validation loss. The code therefore recomputes training loss in eval mode for each checkpoint. The gap is descriptive; an underfit model can have a tiny gap and poor validation loss. Select by the declared validation objective, not smallest gap. The global matrix norm is $\sqrt{\sum_i\|W_i\|_F^2}$, not $\sum_i\|W_i\|_F$.

For a separate MNIST experiment, use the official training partition for training/validation and keep its test partition untouched, change input width to 784 and flatten the images. That larger data experiment was not run. The tested bundled-digits ablation does not certify the original numerical MNIST claims.

Exact normalization and dropout checks

This second standalone CPU program verifies axis formulas, the BatchNorm running-state convention, train/eval distinctions and dropout's distribution. It also checks plain-SGD L2 equivalence and rejects extending that equivalence to momentum: with the same constant data gradient, the coupled penalty enters the velocity, while separate shrinkage does not. The moment expectation is derived above; the sampled tolerance test is only a finite statistical check.

import torch
from torch import nn
torch.set_num_threads(1); torch.manual_seed(40)
x=torch.tensor([[1.,2.],[3.,6.],[5.,10.]],dtype=torch.float64)
bn=nn.BatchNorm1d(2,affine=False,momentum=1.).double()
y=bn(x)
assert torch.allclose(y,(x-x.mean(0))/torch.sqrt(x.var(0,unbiased=False)+bn.eps))
assert torch.allclose(bn.running_mean,x.mean(0))
assert torch.allclose(bn.running_var,x.var(0,unbiased=True))
bn.eval(); before=bn.running_mean.clone()
assert torch.allclose(bn(x),(x-x.mean(0))/torch.sqrt(x.var(0,unbiased=True)+bn.eps))
assert torch.equal(before,bn.running_mean)
ln=nn.LayerNorm(2,elementwise_affine=False).double()
assert torch.allclose(ln(x),(x-x.mean(-1,keepdim=True))/torch.sqrt(x.var(-1,unbiased=False,keepdim=True)+ln.eps))
rms=nn.RMSNorm(2,eps=1e-5,elementwise_affine=False).double()
assert torch.allclose(rms(x),x/torch.sqrt(x.square().mean(-1,keepdim=True)+1e-5))
g=torch.arange(16,dtype=torch.float64).reshape(2,4,2)
gn=nn.GroupNorm(2,4,affine=False).double(); reshaped=g.reshape(2,2,-1)
expected=(reshaped-reshaped.mean(-1,keepdim=True))/torch.sqrt(reshaped.var(-1,unbiased=False,keepdim=True)+gn.eps)
assert torch.allclose(gn(g),expected.reshape_as(g))
d=nn.Dropout(.25); z=torch.full((100000,),2.)
sample=d(z)
assert abs(sample.mean().item()-2)<.02
assert abs(sample.var(unbiased=False).item()-4/3)<.04
assert set(sample.unique().tolist()) <= {0.,float(torch.tensor(8/3))}
d.eval(); assert torch.equal(d(z),z)
assert torch.equal(nn.Dropout(1.)(z),torch.zeros_like(z))
# Plain SGD L2 equivalence; momentum is a negative control, not the same recipe.
w0=torch.tensor([2.,-1.]); grad=torch.tensor([.3,-.4]); lr=.1; decay=.2
w=w0.clone().requires_grad_(); sgd=torch.optim.SGD([w],lr=lr,weight_decay=decay)
w.grad=grad.clone(); sgd.step()
assert torch.allclose(w,(1-lr*decay)*w0-lr*grad)
coupled=w0.clone().requires_grad_()
sgd_m=torch.optim.SGD([coupled],lr=lr,momentum=.9,weight_decay=decay)
decoupled=w0.clone(); velocity=torch.zeros_like(w0)
for _ in range(2):
    coupled.grad=grad.clone(); sgd_m.step()
    velocity=.9*velocity+grad
    decoupled=(1-lr*decay)*decoupled-lr*velocity
assert not torch.allclose(coupled,decoupled)
# AdamW shrink with an explicit zero gradient and no earlier moment state.
w=torch.tensor([2.],requires_grad=True); opt=torch.optim.AdamW([w],lr=.1,weight_decay=.2)
w.grad=torch.zeros_like(w); opt.step()
assert torch.allclose(w,torch.tensor([1.96]))
print('normalization axes, BN state, dropout moments, SGD equivalence, momentum negative control and AdamW shrink passed')

Strength-sweep practice with an answer

Exercise: Extend epochs and sweep dropout or weight decay. Must validation loss form a U, and must 0.5 decay give a norm ten times smaller than 0.05?

Answer: No. Underfitting can emerge with excessive regularization, but finite runs need not be monotonic or U-shaped. The update depends on gradients, rate, parameter groups and step count; norm ratios do not scale inversely with a decay setting by theorem. Record both validation loss and matched eval-mode training loss, checkpoint choice, norms and seeds. The isolated decay variants above make the comparison possible without assuming the result. If extending the search after inspecting test loss, obtain a new independent assessment rather than reusing that test to claim success.

(d) Production reality · 15 min

A model with low training error and a validation gap deserves investigation, not an automatic dropout increase. Check independent sampling units, label delays, feature availability, class balance and domain shift. A stable gap does not prove absence of overfitting; rising validation loss while training loss falls is one symptom, not the definition of every generalization failure.

For normalization, inspect the actual reduction axes, padding and running statistics. BatchNorm can couple examples and, in a sequence layout, positions that should not share future information. LayerNorm over features avoids that particular batch coupling, but the full model still needs causal masking where appropriate. RMSNorm changes the function, not merely its speed. Architecture/version-specific conventions should guide a baseline, followed by validation—not unsupported accounts of a company's failed BN experiment.

Decoupled decay is not guaranteed to reduce final data loss by a fixed percentage. Preserve its optimizer state and parameter-group configuration during resumption. A growing production gap can reflect a changed data-generating process; retraining or adding decay without diagnosing that change may leave it unchanged.

Where this shows up in the rest of the plan

S102 CNNs uses spatial/channel layouts; sequence models often choose feature-wise normalization. Adapter fine-tuning changes trainable parameterization and may still use dropout/decay; it is not simply “regularization instead of dropout.” Evaluation-set overfitting belongs to S095's selection boundary regardless of regularizer. A capstone should justify its ablations, chosen state and untouched final assessment.

(e) Recall + stretch · 10 min

1. Dropout ensemble? Many shared-weight masked subnetworks; ordinary eval approximates rather than exactly averages their predictions.
2. Why identity at eval? Inverted training scaling preserves each activation's expectation for $p<1$; nonlinear network expectations need not match.
3. BN formula? Subtract group mean, divide by square root of biased group variance plus epsilon, apply learned affine parameters; default eval uses running state, whose variance update is unbiased in PyTorch.
4. Why LayerNorm in many transformers? Feature-wise current-input statistics avoid batch dependence and running averages; the actual axes are decisive.
5. Adam-L2 versus AdamW? Coupled penalty enters moment estimates; decoupled shrink does not. Both are defined procedures.
6. Dropout hurts quality? Excessive noise/underfitting is one hypothesis. Compare controlled validation runs and data/mode contracts before concluding.
7. BN and ReLU order? Linear/Conv→BN→ReLU is one common block; pre-activation and other designs exist. Follow the chosen architecture's contract.
8. BN/dropout interaction? Changed second moments can mismatch training and evaluation statistics; downstream placement matters and requires an ablation.
9. Loss 0.01 versus 0.03? A scalar gap alone cannot diagnose overfitting or establish that adding dropout helps. Examine comparable-mode curves, support and shifts.
10. Stretch norm? Compute square root of the sum of squared entries across selected matrices. Compare decay settings without asserting a predetermined norm ratio or quality winner.

Explain-out-loud test

Derive dropout's mean and variance, name the exact BN/LN/GN axes for a tensor shape, and algebraically separate plain-SGD L2 from AdamW. Explain why minimizing a train/validation gap is not the same as maximizing generalization.

What comes next

Convolutional networks add spatial weight sharing; normalization and stochastic regularization still require the same explicit train/eval contracts.

Sources and execution boundary

Dropout primary paper, LayerNorm primary paper, AdamW paper, and PyTorch 2.9 tagged dropout, BatchNorm, normalization and Adam definitions. Both programs are CPU-executed; no MNIST, accelerator, distributed-normalization or production-quality benchmark is claimed.