Regularization: dropout, weight decay, label smoothing, augmentation, early stopping
Before this lesson: lesson 02, training a network (softmax, cross-entropy, SGD), lesson 03, backpropagation and lesson 05, normalization. From other tracks: L2 regularization and the bias–variance trade-off (ML fundamentals track), Adam (optimization track), and the SVD (linear algebra lesson 08) for the last section. Time: about 55 minutes, plus the notebook.
Why this matters #
- LLM pretraining uses AdamW with weight decay 0.1 (GPT-3, LLaMA). Using
Adam(weight_decay=...)instead is a common bug: it is L2, not weight decay. - The original Transformer used dropout 0.1 and label smoothing 0.1.
- Vision transformers trained on ImageNet only (DeiT) need heavy augmentation, label smoothing and stochastic depth; without them they overfit.
- Small scientific datasets rely on early stopping on a validation set, and on augmentations that respect the physics of the data.
1. Predict first #
2. The picture #
Dropout, simplest case #
Take one unit with value and a drop probability . At training time, keep the unit with probability and multiply it by ; otherwise set it to 0. The expected output is
At eval time, use itself. So the network sees the same mean in both modes, but during training it cannot rely on any single unit: any unit may be missing.
Press One training step a few times, then Sample 1000 masks. Each step drops some units (magenta) and scales the kept ones (blue) by ; the violet lines are the weights. The histogram shows the output over many masks. Its mean (blue line) meets the eval output (aqua line). Move : the spread grows like , but the mean stays.
Weight decay, simplest case #
Weight decay multiplies every weight by at every step ( the learning rate, the decay rate). Weights that the loss does not need decay towards 0. For SGD this is the same as adding to the loss (L2 regularization). Adam divides each gradient coordinate by its own running size. If the penalty's gradient goes through Adam, it is divided too, and the shrinkage is no longer the same for every weight. Section 3 shows this in a figure.
Label smoothing #
With label smoothing, the target is not one-hot. The true class gets a bit less than 1, and the other classes get a bit more than 0. The model is never asked to be 100% sure, so the logits stay finite.
Move the logit gap and watch the dot on the loss curve. The left chart compares the smoothed target with the prediction . With smoothing (aqua curve), the loss has a minimum at a finite gap (the dashed lines); beyond it the gradient pushes the gap back down. Without smoothing (grey) the loss keeps falling as . Set to see that case.
Augmentation and early stopping #
If the label does not change when an input is rotated, show the model rotated inputs: the invariance enters through the data. Early stopping limits how far the weights travel from their small initial values, which acts like L2 (section 3, last part).
A unit has value and dropout uses (PyTorch convention). Which values can the training output take?
In PyTorch, is the probability to drop. Kept units are scaled by , so , with probability . The mean is . The first option forgets the scaling; then the mean would be and eval mode would see larger activations than training.
3. Definitions and main results #
A regularizer is a change to the training procedure (the loss, the update rule, the data or the stopping time) that aims to lower the test error, possibly at the cost of a higher training error. Lesson 17 discusses why large networks generalize at all.
Inverted dropout #
With drop probability , draw independently for each entry, and set
- Expectation: . Variance: .
- Backward: , with the same mask as in the forward pass.
- Why "inverted". The original dropout (Srivastava et al., 2014) did not scale during training and multiplied the weights by the keep probability at test time. Scaling during training leaves the eval network unchanged, so nothing special happens at inference.
- Ensemble view. Each mask selects one of thinned subnetworks that share weights. Eval mode approximates the average of their predictions. When everything after the dropout layer is linear, the average is exact (this is the figure above).
Theorem (dropout on the inputs of linear regression is a ridge penalty). Let with independent masks. Then
Proof.
- , so .
- , and the masks are independent, so .
- For any random : . Sum over .
So dropout acts like a ridge penalty whose strength for grows with the energy of feature .
Dropout before BatchNorm. In training, dropout inflates the second moment: . In eval mode it is . A BatchNorm layer after dropout learns running statistics from the inflated training activations, so they do not match the eval activations. Li et al. (2019) call this the variance shift. Place dropout after the last BatchNorm, or avoid mixing the two.
L2 regularization and weight decay #
Theorem (they are the same for SGD). SGD with learning rate on makes exactly the same steps as SGD on with weight decay, .
Proof.
- The gradient of the penalized loss is .
- One step: .
The decay per step is : the best L2 coefficient depends on the learning rate.
Adam with L2. Adam keeps running averages of the gradient, , and of its square, , corrects their start-up bias (), and steps . With L2, the term is added to the gradient before these averages. So the decay moves coordinate by about :
- weights with large past gradients (large ) are decayed less;
- weights with small gradients are decayed more.
Loshchilov and Hutter (2019, Proposition 2) prove that, for any preconditioner that is not a multiple of , no L2 coefficient makes the method equal to weight decay.
AdamW decouples the decay from the gradient step: the moments use only, and
with base learning rate and schedule multiplier (their Algorithm 2). Every weight now
shrinks at the same rate. PyTorch's AdamW first applies
and then the Adam step, so its decay per step is lr × weight_decay and follows the learning-rate
schedule. PyTorch's Adam(weight_decay=...) is L2, not AdamW.
Change the number of steps and the gradient size . The chart shows how much the decay term alone shrank each weight (the run with decay minus the run without). With AdamW, both weights shrink by about (the two aqua lines almost overlap). With Adam + L2, the weight with no gradient (solid) is pulled to 0 within about 100 steps, while the weight with large gradients (dashed) shrinks about 10 times less than under AdamW.
With normalization layers. A weight followed by BatchNorm or LayerNorm is scale-invariant (lesson 05): decay does not change the function, but it keeps small and so keeps the effective learning rate up. Biases and normalization gains are usually excluded from decay.
You train with SGD, learning rate , and add with to the loss. By which factor does the penalty multiply the weights at each step?
The penalty adds to the gradient, and the step multiplies the gradient by : . The factor in the penalty cancels with the 2 from the derivative. This is why the decay per step depends on the learning rate.
Label smoothing #
With classes, true class and smoothing (PyTorch convention):
So the true class gets and every other class gets . (Some texts give the true class and the others .)
Gradient. Since ,
because . This is the same form as plain cross-entropy, with in place of the one-hot target.
Minimal loss and bounded logits. For a fixed , , where is the entropy of and (Gibbs' inequality). So the loss is minimal exactly at , and it never goes below . At , the logit gap between the true class and any other class is
which is finite. Without smoothing the optimum needs an infinite gap, so the logits keep growing and the model becomes overconfident. Müller et al. (2019) found that label smoothing improves calibration, but makes the network a worse teacher for knowledge distillation.
classes and label smoothing . What is the optimal logit gap ?
The target is and . The optimum is , so the gap is . The first option uses , which mixes up " spread over the other classes" with " per class". An infinite gap is the case .
Data augmentation encodes invariances #
Suppose the label does not change under a finite group of transformations (rotations by 90°, flips, cyclic shifts): for every . With drawn uniformly from , data augmentation minimizes
For a loss that is convex in the prediction, the averaged predictor is -invariant, and by Jensen's inequality . Since is invariant, : the augmented objective prefers invariant predictors.
Additive input noise is another simple case. For linear regression on samples, with independent noise added to each input ,
which is ridge again (Bishop, 1995). An equivariant architecture (a CNN, or an E(3)-equivariant graph network) builds the same symmetry into the model instead of the data (lessons 07 and 13).
An augmentation must keep the label and the physics. Rotating a handwritten "6" by 180° gives a "9". Flipping a regional weather map north–south makes cyclones turn the wrong way for their hemisphere. Rescaling a seismic waveform changes the magnitude it encodes.
Early stopping acts like ridge #
Take least squares, , with gradient , where and . Assume is invertible. Write (its eigendecomposition; the are the squared singular values of ) and . Ridge with strength minimizes , so it solves .
Theorem (filter factors). Gradient descent from with step gives, along each eigenvector with eigenvalue , the component after steps. Ridge with strength gives .
Proof.
- , because .
- Hence , that is, .
- Ridge solves , so .
- In the eigenbasis both matrices are diagonal, which gives the two filter factors.
Reading the factors. For both factors are about 1: these directions are fitted. For , GD gives about (because ) and ridge gives about . They match when
So stopping early suppresses the directions with small eigenvalues, which are the ones most affected by noise, just as ridge does. The stopping time is chosen on a validation set, never on the test set.
Move the number of steps . The violet path is gradient descent from 0 (the dots are steps); the aqua dashed curve is the ridge solution for all , and the aqua dot is ridge with . Both first move along the steep direction of the loss (large eigenvalue ) and reach along the flat direction () only later. The right chart shows the two filter factors; the dots mark and .
Answers to section 1. (1) No. The minimum is with : . (2) With Adam + L2, the weight with no gradient shrinks much faster: its is tiny, so the decay term is divided by a tiny number. With AdamW both shrink at the same rate, per step.
4. Worked example #
(a) Dropout. , , mask draw .
- Training output: .
- Over many draws: and .
- Backward with this mask: .
- Eval: .
(b) Label smoothing. , , true class 0, logits .
- Target: .
- Prediction: .
- Loss: .
- Gradient: . It still pushes the gap up, because the optimal gap is . Above 2.89 the gradient pushes back.
- Without smoothing, the gradient is , and its first entry is negative for every gap.
(c) Adam + L2 vs AdamW. Learning rate 0.01, , 10 steps, two weights at 1. The data gradients alternate in sign, with size 0 (first weight) or 10 (second weight). Adam moves the second weight even without decay (by about 0.015 here), so we compare each run with the same run without decay. The shrinkage caused by the decay:
| no-gradient weight | large-gradient weight | |
|---|---|---|
| Adam + L2 | 0.0997 | 0.00009 |
| AdamW | 0.0010 | 0.0010 |
AdamW shrinks both by about . Adam + L2 shrinks the no-gradient weight 100 times more, and the large-gradient weight 11 times less.
(d) Early stopping. and , so .
| eigenvalue | GD: | ridge: |
|---|---|---|
| 1 | ||
| 0.1 | ||
| 0.01 |
Both keep the large- directions and shrink the small- ones; GD's filter is sharper.
Gradient descent with is stopped after steps on a least-squares problem. A direction has eigenvalue . About how much of along it has GD fitted?
The filter factor is . Here , so the factor is close to . The matching ridge strength is , not , which would give .
5. In code #
PyTorch checks the worked examples (a)–(c):
import torch
import torch.nn.functional as F
torch.manual_seed(0)
x = torch.tensor([1., 2., 3., 4.], dtype=torch.float64).repeat(200_000, 1) # 200k copies
y = F.dropout(x, p=0.5, training=True)
print(sorted(y[:, 3].unique().tolist())) # [0.0, 8.0]: kept entries scaled by 1/(1-p)
print(round((y == 0).double().mean().item(), 3)) # 0.5: fraction dropped
print(y.mean(0).numpy().round(1)) # [1. 2. 3. 4.] E[y] = x
print(y.var(0).numpy().round(1)) # [ 1. 4. 9. 16.] Var = x^2 p/(1-p)
print(torch.equal(F.dropout(x, p=0.5, training=False), x)) # True: identity at eval
z = torch.tensor([[2.0, 0.0, 0.0]], dtype=torch.float64, requires_grad=True)
loss = F.cross_entropy(z, torch.tensor([0]), label_smoothing=0.15)
loss.backward()
print(round(loss.item(), 4), z.grad.numpy().round(4)) # 0.4395 [[-0.113 0.0565 0.0565]]
def run(opt_cls, weight_decay, steps=10):
"""Two weights start at 1. Their data gradients alternate in sign, with sizes 0 and 10."""
theta = torch.ones(2, dtype=torch.float64, requires_grad=True)
size = torch.tensor([0.0, 10.0], dtype=torch.float64)
opt = opt_cls([theta], lr=0.01, weight_decay=weight_decay)
for t in range(steps):
theta.grad = size * (-1.0) ** t # the data gradient of this step
opt.step()
return theta.detach()
for name, cls in [("Adam + L2", torch.optim.Adam), ("AdamW", torch.optim.AdamW)]:
shrink = run(cls, 0.0) - run(cls, 0.01) # the part of the movement caused by the decay
print(f"{name:9s} no-gradient weight {shrink[0]:.5f} large-gradient weight {shrink[1]:.5f}")
# Adam + L2 no-gradient weight 0.09965 large-gradient weight 0.00009
# AdamW no-gradient weight 0.00100 large-gradient weight 0.00099
Early stopping against ridge on a least-squares problem with a wide spread of eigenvalues:
import numpy as np
rng = np.random.default_rng(0)
n, d = 200, 20
X = rng.standard_normal((n, d)) * np.logspace(0, -1.5, d) # feature scales from 1 down to 0.03
y = X @ rng.standard_normal(d) + 0.1 * rng.standard_normal(n)
S, b = X.T @ X / n, X.T @ y / n # the loss (1/2n)|Xw - y|^2 has gradient S w - b
eta, t = 0.5, 400
w = np.zeros(d)
for _ in range(t): # gradient descent from 0, stopped after t steps
w -= eta * (S @ w - b)
lam = 1 / (eta * t) # predicted equivalent ridge strength: 0.005
w_ridge = np.linalg.solve(S + lam * np.eye(d), b)
w_ls = np.linalg.solve(S, b)
s, V = np.linalg.eigh(S) # filter factor = component of w / component of w_ls
f_gd, f_ridge = (V.T @ w) / (V.T @ w_ls), (V.T @ w_ridge) / (V.T @ w_ls)
print(np.allclose(f_gd, 1 - (1 - eta * s) ** t), np.allclose(f_ridge, s / (s + lam))) # True True
for i in (0, 5, 10, 19):
print(f"s={s[i]:.4f} GD {f_gd[i]:.3f} ridge {f_ridge[i]:.3f}")
# s=0.0009 GD 0.158 ridge 0.147
# s=0.0070 GD 0.755 ridge 0.584
# s=0.0332 GD 0.999 ridge 0.869
# s=0.9810 GD 1.000 ridge 0.995
6. Where it shows up #
- LLM pretraining. AdamW with weight decay 0.1 (GPT-3, LLaMA), no decay on biases and norm gains, and usually no dropout: with one pass over a huge dataset, overfitting is not the main risk. Dropout returns when fine-tuning on small data.
- Vision. DeiT trains ViTs on ImageNet-1k only, with strong augmentation (RandAugment, Mixup, CutMix), label smoothing 0.1 and stochastic depth (randomly dropping whole residual branches).
- ML for science. Labelled data are often small, so early stopping on a held-out set is standard. Split the held-out set by time or region, not at random, or neighbouring samples leak. Augment only with true symmetries: rotations are fine for molecules, but not a north–south flip of a regional weather map. Learned simulators add noise to their inputs during training so that long autoregressive rollouts stay stable (Sanchez-Gonzalez et al., 2020).
- Uncertainty. Monte Carlo dropout keeps dropout on at test time and averages several predictions; their spread is a cheap uncertainty estimate (Gal and Ghahramani, 2016).
7. Common mistakes #
Adam(weight_decay=λ)is L2, not AdamW. UseAdamWfor decoupled decay.- Tuning without the learning rate. In PyTorch's AdamW the decay per step is lr × λ, so changing lr changes the decay, and the decay fades with a cosine schedule.
- Decaying biases and normalization gains. They are usually excluded; decaying pulls the output scale of every normalized layer towards 0.
- Forgetting
model.eval(). Dropout stays on and the predictions are noisy. - Wrong meaning of . PyTorch's is the drop probability; Srivastava et al. and TensorFlow 1's
keep_probuse the keep probability. - Dropout before BatchNorm. Dropout inflates by during training, so BN's running variance does not match the eval activations.
- Label-changing augmentation. Rotating a "6" by 180° gives a "9"; flipping a regional weather map breaks the physics.
- Early stopping on the test set. That leaks the test set into model selection; use a validation set.
8. Check yourself #
- Dropout with on a unit with value . Which values can take, with which
probabilities? What are and ?
Answer
with probability 0.75, and with probability 0.25. and . - In PyTorch, what does a dropout layer do to its input in eval mode, and why is no rescaling needed there?
Answer
It returns the input unchanged. Inverted dropout already divides by during training, so the expected training output equals the eval output. - , : give the smoothed target, the minimal loss and the optimal logit gap.
Answer
True class , others . Minimal loss . Gap . - PyTorch
AdamW(lr=3e-4, weight_decay=0.1)with a constant learning rate: by which factor does the decay alone shrink a weight over 10,000 steps?Answer
The factor per step is , so . - Linear regression with dropout on the inputs, and standardized features (). Which ridge penalty does dropout add to the summed squared loss?
Answer
. For , , which is ridge strength 1 on the mean squared loss. - Gradient descent with is stopped after steps. Which ridge strength does it resemble, and which eigen-directions of does it fit?
Answer
. Directions with are fitted almost fully; directions with are shrunk to about of their value.
Write inverted dropout with drop probability (training and eval).
Training: with . Eval: .
Why does inverted dropout divide by ?
So that : the training output has the same mean as the eval output, and eval needs no change.
Dropout with on a unit with value : what is ?
.
Show that SGD with the L2 penalty is weight decay.
The step is .
Why is L2 regularization not weight decay under Adam?
The L2 gradient is divided by like the rest of the gradient, so weights with large past gradients are decayed less and weights with small gradients more.
What is the AdamW update in PyTorch?
, then the usual Adam step computed from only.
What is the label-smoothed target with classes and smoothing ?
: the true class gets , every other class .
What is the minimal cross-entropy under label smoothing, and where is it reached?
The entropy , reached at , because . For , it is about .
Which ridge strength does early stopping after GD steps with step resemble?
: the filter factors and agree for large and for small eigenvalues .
What does data augmentation with a symmetry group of the labels encourage?
Predictors that are invariant under : the averaged predictor has a lower or equal augmented loss (Jensen, for a convex loss).
9. Interview angle #
- "Why scale by in dropout?" — So that , and nothing changes at eval time.
- "Weight decay vs L2?" — Equal for SGD (factor per step); different for Adam, because the L2 gradient is divided by . AdamW decouples the decay.
- "What does label smoothing do to the logits?" — The optimum is with a finite gap ; the loss never reaches 0, and the model is less overconfident.
- "Is early stopping a regularizer?" — Yes: for least squares it acts like ridge with .
- "Why do LLMs often train without dropout?" — One epoch over a huge dataset; the regularization comes from data scale and weight decay.
10. Practice #
The notebook for this lesson runs in Colab. Each exercise has a check cell and a folded solution.
- Inverted dropout (forward and backward with the same mask), with a check of the mean and the variance.
- Adam + L2 and AdamW written by hand on two weights, reproducing the numbers of section 4 and checked against
torch.optim. - Label-smoothed cross-entropy and its gradient, checked against
F.cross_entropy(..., label_smoothing=). - Early-stopping and ridge filter factors on a random least-squares problem.
- A small MLP on the scikit-learn digits: validation accuracy with and without dropout and weight decay.
11. Watch and read #
- Simon Prince, Understanding Deep Learning, free PDF: chapter 9 (§9.1.2 L2 regularization, §9.3.1 early stopping, §9.3.3 dropout, §9.3.4 applying noise, §9.3.8 augmentation).
- d2l.ai, §5.6 Dropout.
- Srivastava et al., Dropout: A Simple Way to Prevent Neural Networks from Overfitting (JMLR, 2014).
- Loshchilov and Hutter, Decoupled Weight Decay Regularization (2019), §2 (Propositions 1–3, Algorithm 2).
- Müller, Kornblith and Hinton, When Does Label Smoothing Help? (2019).
- Zhang et al., Understanding Deep Learning Requires Rethinking Generalization (2017): networks fit random labels.
- Li et al., Understanding the Disharmony between Dropout and Batch Normalization by Variance Shift (2019).
Runs in your browser with nothing to install. The notebook repeats the key demos and has exercises with checks; solutions are inside, folded.