Unit 05.04: Dropout, weight decay, and early stopping

Three ways to stop a model memorising. They work differently and are worth understanding separately before you combine them.

Three different brakes

Dropout randomly zeroes activations during training, so no single unit can be relied on. Weight decay penalises large weights, keeping the function smoother. Early stopping simply stops when validation stops improving.

The data here is pure noise, so any drop in training loss is memorisation by definition:

import torch
from torch import nn

torch.manual_seed(0)
X, y = torch.randn(120, 8), torch.randn(120, 1)     # noise: nothing to learn
Xtr, ytr, Xva, yva = X[:80], y[:80], X[80:], y[80:]


def train(dropout=0.0, weight_decay=0.0, steps=400):
    torch.manual_seed(0)
    model = nn.Sequential(nn.Linear(8, 128), nn.ReLU(), nn.Dropout(dropout), nn.Linear(128, 1))
    opt = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=weight_decay)
    best_val, best_step = float("inf"), 0
    for step in range(steps):
        model.train(); opt.zero_grad()
        nn.MSELoss()(model(Xtr), ytr).backward()
        opt.step()
        model.eval()
        with torch.no_grad():
            val = nn.MSELoss()(model(Xva), yva).item()
        if val < best_val:
            best_val, best_step = val, step
    model.eval()
    with torch.no_grad():
        return nn.MSELoss()(model(Xtr), ytr).item(), val, best_val, best_step


for label, kwargs in [("none", {}), ("dropout 0.5", {"dropout": 0.5}),
                      ("weight decay 0.1", {"weight_decay": 0.1})]:
    tr, final_val, best_val, best_step = train(**kwargs)
    print(f"{label:18} train {tr:.4f}  val(final) {final_val:.4f}  "
          f"val(best) {best_val:.4f} at step {best_step}")

# Training loss near zero on random targets is memorisation. Note how much
# earlier the best validation score occurs than the last step -- that gap is
# the argument for early stopping.

The unregularised model drives training loss towards zero on random targets — textbook memorisation. Both regularisers raise that training loss, which is the intended effect.

Now look at the gap between val(best) and the step it occurred on. Best validation arrives long before the final step, and everything after it is wasted computation making the model worse.

The mistake this prevents

Adding dropout and weight decay together, then tuning both at once. When the result changes you cannot attribute it. Change one thing at a time.

Takeaway

Regularisation should raise training loss. If it does not, it is not doing anything.

Mark complete