Dropout and Normalization: Two Layers Every Diagram Has

Initialization and Optimizers chose where training starts and how it steps. This post covers the two kinds of layer that get inserted into a network to make training behave, and that show up in nearly every architecture diagram:

• Dropout randomly switches off units during training. It’s a regularizer, the subject of Features and Overfitting, done by noise instead of a penalty.
• Normalization re-standardizes activations inside the network, on every forward pass. It’s the feature scaling of Learning Is Optimization, moved inside and applied continuously, and it changes what the gradient and the learning rate mean.

Both are tested on the router from Initialization and Optimizers, with the series’ usual methods. One of them turns out to matter much more for language models than the other. Language models are trained with normalization in every block, and many are pretrained with no dropout at all.

Table of contents

Dropout: training a crowd of thinner networks

Dropout (Hinton et al., 2012; Srivastava et al., 2014) does one thing during training: every hidden unit’s output is set to zero with probability $p$, independently, on every step. (The original also dropped some inputs. In this post $p$ is the dropout rate, a local symbol, not the model probability $p_\theta$ of AI Engineering, from the Inside.) A different random set of units is missing on each step, so no unit can count on any particular other unit being there. The intended effect is that units stop co-adapting, where one unit’s feature only works because of another’s quirks, and each learns something useful on its own.

At test time nothing is dropped. That creates a mismatch: during training a unit’s input was, on average, a fraction $1 - p$ of what it will be at test time. Inverted dropout, what every library implements, fixes it during training: keep each unit with probability $1 - p$ and divide the kept ones by $1 - p$. A unit with output $a$ then contributes

on average, the same as at test time, so the test-time network needs no change. The backward pass is just as simple: the gradient flows through the same mask, scaled by the same $1/(1-p)$, and dropped units get none.

Why it regularizes. A network with $n$ hidden units has $2^n$ ways to drop some of them, so training with dropout trains a huge family of thinner networks that all share weights. The test-time network, with every unit on and the $1/(1-p)$ scaling built in, is a cheap stand-in for averaging that whole family’s predictions. For one linear layer it’s exact: the full network’s pre-activations are the probability-weighted average of all the thinned networks’. On a softmax output it gives a geometric mean of their predictions, and through deeper nonlinearities it’s an approximation, one that works well in practice. It’s the bias–variance trade of Features and Overfitting: each thinned network is noisier, and averaging them lowers the variance.

Here it is in net.py, the network code this post uses. The dropout lines are in forward and loss_and_grads; the file also holds init and fit, the mini-batch training loop:

def forward(layers, X, norm=None, drop=0.0, rng=None):
    """Each hidden layer: linear, then (optionally) a norm, ReLU, then dropout."""
    h, caches = X, []
    for i, L in enumerate(layers):
        z = h @ L["W"] + L["b"]
        if i == len(layers) - 1:
            caches.append((h, None, None, None))
            return z, caches
        nc = None
        if norm == "layer":
            z, nc = layernorm(z, L["g"], L["beta"])
        elif norm == "rms":
            z, nc = rmsnorm(z, L["g"])
        a = np.maximum(z, 0)
        mask = None
        if drop and rng is not None:                       # inverted dropout, training only
            mask = (rng.random(a.shape) >= drop) / (1 - drop)
            a = a * mask
        caches.append((h, z, nc, mask))
        h = a

def loss_and_grads(layers, X, y, norm=None, drop=0.0, rng=None):
    logits, caches = forward(layers, X, norm, drop, rng)
    g = (softmax(logits) - np.eye(logits.shape[1])[y]) / len(y)      # p - y
    grads = [None] * len(layers)
    for i in reversed(range(len(layers))):
        h, z, nc, mask = caches[i]
        if i < len(layers) - 1:                            # back through dropout, ReLU, norm
            if mask is not None:
                g = g * mask
            g = g * (z > 0)
            G = {}
            if norm == "layer":
                g, G["g"], G["beta"] = layernorm_back(g, layers[i]["g"], nc)
            elif norm == "rms":
                g, G["g"] = rmsnorm_back(g, layers[i]["g"], nc)
        else:
            G = {}
        G.update({"W": h.T @ g, "b": g.sum(0)})
        grads[i] = G
        g = g @ layers[i]["W"].T
    return cross_entropy(logits, y), grads

Full script: net.py

The backward pass, with and without the norm layers below, checks out against finite differences (grad_check.py):

import numpy as np
from net import init, loss_and_grads

# The whole network's backward pass, with each norm, against finite differences.
rng = np.random.default_rng(0)
X, y = rng.normal(size=(6, 5)), rng.integers(0, 4, 6)
for norm in (None, "layer", "rms"):
    layers, worst = init([5, 7, 7, 4], rng, norm=norm), 0.0
    _, grads = loss_and_grads(layers, X, y, norm)
    for L, G in zip(layers, grads):
        for k in L:
            for i in np.ndindex(L[k].shape):
                old = L[k][i]
                L[k][i] = old + 1e-6; up = loss_and_grads(layers, X, y, norm)[0]
                L[k][i] = old - 1e-6; down = loss_and_grads(layers, X, y, norm)[0]
                L[k][i] = old
                num = (up - down) / 2e-6
                worst = max(worst, abs(num - G[k][i]) / max(1e-12, abs(num) + abs(G[k][i])))
    print(f"norm={norm}: largest relative error {worst:.1e}")

Full script: grad_check.py

norm=None: largest relative error 3.8e-08
norm=layer: largest relative error 6.7e-08
norm=rms: largest relative error 4.9e-07

Does dropout help the router? The test: a network with two hidden layers of 64, big enough to memorize its 120 training questions, trained with Adam. Each combination of dropout rate and AdamW weight decay is scored on validation, and the best one on test. The loop at the heart of dropout_grid.py:

STEPS = int(sys.argv[1]) if len(sys.argv) > 1 else 3000
CONFIGS = [(p, lam) for lam in (0.0, 0.3, 1.0, 3.0, 10.0) for p in (0.0, 0.1, 0.3, 0.5, 0.7)]
rows = {c: [] for c in CONFIGS}
for seed in range(20):
    ((Xtr, ytr), (Xva, yva), (Xte, yte)), _ = four_way(*split(seed))
    for p, lam in CONFIGS:
        rng = np.random.default_rng(seed)
        layers = init([Xtr.shape[1], 64, 64, 4], rng)
        fit(layers, Adam(1e-3, decay=lam), Xtr, ytr, STEPS, rng, drop=p)
        zv, zt = forward(layers, Xva)[0], forward(layers, Xte)[0]      # no dropout at test time
        top2 = np.sort(zv, 1)[:, -2:]                              # the two largest logits
        wrong = zv.argmax(1) != yva
        rows[(p, lam)].append((cross_entropy(zv, yva), (zv.argmax(1) == yva).mean(),
                               softmax(zv).max(1).mean(), cross_entropy(zt, yte),
                               (top2[:, 1] - top2[:, 0]).mean(), softmax(zv).max(1)[wrong].mean()))

Full script: dropout_grid.py

3,000 steps of Adam. Differences are paired, against no dropout at the same decay.
dropout 0.0  decay 0.0  val CE 1.115                  val acc 0.667                  mean top p 0.857  top-2 logit gap  3.85  top p when wrong 0.752
dropout 0.1  decay 0.0  val CE 1.228  +0.113 ± 0.052  val acc 0.680  +0.013 ± 0.023  mean top p 0.874  top-2 logit gap  4.43  top p when wrong 0.786
dropout 0.3  decay 0.0  val CE 1.462  +0.347 ± 0.080  val acc 0.686  +0.019 ± 0.032  mean top p 0.910  top-2 logit gap  5.80  top p when wrong 0.840
dropout 0.5  decay 0.0  val CE 1.661  +0.546 ± 0.132  val acc 0.679  +0.011 ± 0.027  mean top p 0.926  top-2 logit gap  6.93  top p when wrong 0.860
dropout 0.7  decay 0.0  val CE 1.513  +0.397 ± 0.116  val acc 0.684  +0.016 ± 0.029  mean top p 0.913  top-2 logit gap  6.54  top p when wrong 0.843
dropout 0.0  decay 0.3  val CE 0.953                  val acc 0.683                  mean top p 0.807  top-2 logit gap  2.95  top p when wrong 0.711
dropout 0.1  decay 0.3  val CE 0.950  -0.003 ± 0.034  val acc 0.693  +0.010 ± 0.027  mean top p 0.822  top-2 logit gap  3.18  top p when wrong 0.719
dropout 0.3  decay 0.3  val CE 0.974  +0.021 ± 0.044  val acc 0.704  +0.021 ± 0.038  mean top p 0.845  top-2 logit gap  3.58  top p when wrong 0.751
dropout 0.5  decay 0.3  val CE 1.018  +0.065 ± 0.058  val acc 0.704  +0.021 ± 0.040  mean top p 0.864  top-2 logit gap  3.99  top p when wrong 0.772
dropout 0.7  decay 0.3  val CE 1.016  +0.062 ± 0.060  val acc 0.701  +0.019 ± 0.034  mean top p 0.864  top-2 logit gap  4.12  top p when wrong 0.767
dropout 0.0  decay 1.0  val CE 0.879                  val acc 0.680                  mean top p 0.762  top-2 logit gap  2.41  top p when wrong 0.658
dropout 0.1  decay 1.0  val CE 0.869  -0.011 ± 0.046  val acc 0.689  +0.009 ± 0.024  mean top p 0.787  top-2 logit gap  2.67  top p when wrong 0.675
dropout 0.3  decay 1.0  val CE 0.865  -0.014 ± 0.054  val acc 0.698  +0.017 ± 0.041  mean top p 0.792  top-2 logit gap  2.76  top p when wrong 0.683
dropout 0.5  decay 1.0  val CE 0.846  -0.033 ± 0.057  val acc 0.715  +0.035 ± 0.034  mean top p 0.793  top-2 logit gap  2.76  top p when wrong 0.692
dropout 0.7  decay 1.0  val CE 0.809  -0.070 ± 0.056  val acc 0.714  +0.034 ± 0.032  mean top p 0.775  top-2 logit gap  2.55  top p when wrong 0.662
dropout 0.0  decay 3.0  val CE 0.809                  val acc 0.694                  mean top p 0.719  top-2 logit gap  1.94  top p when wrong 0.605
dropout 0.1  decay 3.0  val CE 0.806  -0.003 ± 0.017  val acc 0.690  -0.004 ± 0.020  mean top p 0.715  top-2 logit gap  1.93  top p when wrong 0.596
dropout 0.3  decay 3.0  val CE 0.799  -0.010 ± 0.022  val acc 0.689  -0.005 ± 0.024  mean top p 0.705  top-2 logit gap  1.86  top p when wrong 0.582
dropout 0.5  decay 3.0  val CE 0.781  -0.028 ± 0.025  val acc 0.707  +0.014 ± 0.025  mean top p 0.678  top-2 logit gap  1.68  top p when wrong 0.561
dropout 0.7  decay 3.0  val CE 0.806  -0.004 ± 0.035  val acc 0.698  +0.004 ± 0.025  mean top p 0.606  top-2 logit gap  1.25  top p when wrong 0.492
dropout 0.0  decay 10.0 val CE 0.940                  val acc 0.670                  mean top p 0.482  top-2 logit gap  0.67  top p when wrong 0.411
dropout 0.1  decay 10.0 val CE 0.962  +0.022 ± 0.018  val acc 0.647  -0.022 ± 0.028  mean top p 0.469  top-2 logit gap  0.62  top p when wrong 0.403
dropout 0.3  decay 10.0 val CE 1.021  +0.081 ± 0.027  val acc 0.629  -0.041 ± 0.043  mean top p 0.430  top-2 logit gap  0.46  top p when wrong 0.376
dropout 0.5  decay 10.0 val CE 1.104  +0.164 ± 0.031  val acc 0.620  -0.050 ± 0.042  mean top p 0.380  top-2 logit gap  0.33  top p when wrong 0.342
dropout 0.7  decay 10.0 val CE 1.246  +0.306 ± 0.036  val acc 0.443  -0.227 ± 0.042  mean top p 0.312  top-2 logit gap  0.07  top p when wrong 0.304
chosen on validation: dropout 0.5, decay 3.0; without dropout: decay 3.0
test CE: 0.776 chosen, 0.798 decay only, 1.083 neither. Dropout's own share: -0.022 ± 0.024; decay's: -0.285 ± 0.097

And four times as long, the same script with 12000 as its argument, which dropout_grid_long.py passes. The grid trains 500 networks per run (20 splits × 25 settings) on a laptop’s CPU; on an Apple M4, the 3,000-step grid and the post’s other scripts take about 7 minutes in all, and the run time grows in proportion to the step count:

12,000 steps of Adam. Differences are paired, against no dropout at the same decay.
dropout 0.0  decay 0.0  val CE 1.467                  val acc 0.689                  mean top p 0.903  top-2 logit gap  5.55  top p when wrong 0.833
dropout 0.1  decay 0.0  val CE 1.787  +0.320 ± 0.085  val acc 0.686  -0.003 ± 0.025  mean top p 0.924  top-2 logit gap  7.21  top p when wrong 0.862
dropout 0.3  decay 0.0  val CE 2.516  +1.049 ± 0.173  val acc 0.691  +0.003 ± 0.024  mean top p 0.954  top-2 logit gap 10.68  top p when wrong 0.914
dropout 0.5  decay 0.0  val CE 2.847  +1.380 ± 0.324  val acc 0.706  +0.018 ± 0.022  mean top p 0.960  top-2 logit gap 13.47  top p when wrong 0.925
dropout 0.7  decay 0.0  val CE 3.533  +2.066 ± 0.481  val acc 0.696  +0.008 ± 0.030  mean top p 0.965  top-2 logit gap 16.41  top p when wrong 0.933
dropout 0.0  decay 0.3  val CE 1.294                  val acc 0.654                  mean top p 0.850  top-2 logit gap  4.14  top p when wrong 0.775
dropout 0.1  decay 0.3  val CE 1.230  -0.064 ± 0.098  val acc 0.696  +0.043 ± 0.052  mean top p 0.883  top-2 logit gap  4.84  top p when wrong 0.802
dropout 0.3  decay 0.3  val CE 1.326  +0.031 ± 0.122  val acc 0.688  +0.034 ± 0.048  mean top p 0.892  top-2 logit gap  5.42  top p when wrong 0.805
dropout 0.5  decay 0.3  val CE 1.296  +0.001 ± 0.116  val acc 0.695  +0.041 ± 0.044  mean top p 0.902  top-2 logit gap  5.71  top p when wrong 0.816
dropout 0.7  decay 0.3  val CE 1.241  -0.053 ± 0.131  val acc 0.714  +0.060 ± 0.035  mean top p 0.900  top-2 logit gap  5.71  top p when wrong 0.824
dropout 0.0  decay 1.0  val CE 1.086                  val acc 0.680                  mean top p 0.848  top-2 logit gap  3.74  top p when wrong 0.760
dropout 0.1  decay 1.0  val CE 1.111  +0.025 ± 0.079  val acc 0.691  +0.011 ± 0.028  mean top p 0.854  top-2 logit gap  3.99  top p when wrong 0.767
dropout 0.3  decay 1.0  val CE 1.029  -0.057 ± 0.048  val acc 0.695  +0.015 ± 0.017  mean top p 0.844  top-2 logit gap  3.83  top p when wrong 0.738
dropout 0.5  decay 1.0  val CE 0.919  -0.167 ± 0.058  val acc 0.699  +0.019 ± 0.019  mean top p 0.834  top-2 logit gap  3.48  top p when wrong 0.716
dropout 0.7  decay 1.0  val CE 0.842  -0.244 ± 0.065  val acc 0.713  +0.033 ± 0.022  mean top p 0.808  top-2 logit gap  2.99  top p when wrong 0.697
dropout 0.0  decay 3.0  val CE 1.010                  val acc 0.667                  mean top p 0.806  top-2 logit gap  2.90  top p when wrong 0.710
dropout 0.1  decay 3.0  val CE 0.922  -0.087 ± 0.037  val acc 0.688  +0.020 ± 0.014  mean top p 0.792  top-2 logit gap  2.73  top p when wrong 0.683
dropout 0.3  decay 3.0  val CE 0.814  -0.195 ± 0.045  val acc 0.700  +0.033 ± 0.023  mean top p 0.751  top-2 logit gap  2.28  top p when wrong 0.631
dropout 0.5  decay 3.0  val CE 0.774  -0.236 ± 0.064  val acc 0.706  +0.039 ± 0.028  mean top p 0.703  top-2 logit gap  1.85  top p when wrong 0.579
dropout 0.7  decay 3.0  val CE 0.799  -0.210 ± 0.083  val acc 0.708  +0.040 ± 0.029  mean top p 0.609  top-2 logit gap  1.27  top p when wrong 0.495
dropout 0.0  decay 10.0 val CE 0.958                  val acc 0.663                  mean top p 0.470  top-2 logit gap  0.62  top p when wrong 0.405
dropout 0.1  decay 10.0 val CE 0.957  -0.001 ± 0.021  val acc 0.662  +0.000 ± 0.031  mean top p 0.468  top-2 logit gap  0.58  top p when wrong 0.402
dropout 0.3  decay 10.0 val CE 1.027  +0.069 ± 0.024  val acc 0.650  -0.013 ± 0.038  mean top p 0.425  top-2 logit gap  0.43  top p when wrong 0.375
dropout 0.5  decay 10.0 val CE 1.121  +0.164 ± 0.033  val acc 0.636  -0.026 ± 0.041  mean top p 0.368  top-2 logit gap  0.29  top p when wrong 0.332
dropout 0.7  decay 10.0 val CE 1.255  +0.298 ± 0.039  val acc 0.426  -0.236 ± 0.056  mean top p 0.306  top-2 logit gap  0.05  top p when wrong 0.298
chosen on validation: dropout 0.5, decay 3.0; without dropout: decay 10.0
test CE: 0.782 chosen, 0.973 decay only, 1.368 neither. Dropout's own share: -0.191 ± 0.041; decay's: -0.395 ± 0.167

Read the rows in groups of five, one weight-decay setting at a time. The differences are paired against no dropout at the same decay, so each group isolates what dropout adds. They’re validation numbers, used to make the choice and picked from 25 cells, so the best of them are optimistic; the test lines at the bottom are the honest ones. Each ”±” is a standard error of the paired difference over the 20 splits, corrected for the overlap between splits that share data (the corrected resampled t-test of Nadeau and Bengio: the standard deviation times $\sqrt{1/20 + 40/120}$, with 40 validation and 120 training questions per split). The chosen setting is the one with the lowest mean validation cross-entropy.

• Without weight decay, dropout makes the router worse. Every dropout rate raises the validation cross-entropy: by up to $0.546 \pm 0.132$ at 3,000 steps, and by up to 2.07 at 12,000. Accuracy doesn’t move beyond its error. The last two columns say why. To keep its training answers right while units are missing, the network learns much larger gaps between the winning logit and the runner-up, so the right class still wins through the noise. At test time the noise is gone and the margins stay: at 12,000 steps, dropout 0.5 more than doubles the average top-two gap, from 5.55 to 13.47, and the router is 92.5% confident even on the questions it gets wrong. Dropout on its own puts no limit on those margins.
• Weight decay supplies the limit, and then dropout helps. At decay 3.0 the margins stay small (1.85 at dropout 0.5, 12,000 steps), and dropout now lowers validation cross-entropy, by $0.236 \pm 0.064$ at 12,000 steps. Accuracy moves the same way ($+0.039 \pm 0.028$) but within about 1.4 standard errors, so on accuracy it’s a tie.
• Too much of both underfits. At decay 10 the router’s average confidence falls below 0.5, and at dropout 0.7 its accuracy collapses to 0.43–0.44.
• On test, most of the gain is weight decay’s, until training runs long. The last lines split the chosen model’s improvement in two, against the best model without dropout, also chosen on validation. At 3,000 steps decay does nearly all of it ($-0.285 \pm 0.097$) and dropout’s own share is a tie ($-0.022 \pm 0.024$). At 12,000 steps, when the network has had time to overfit, dropout’s share is real: $-0.191 \pm 0.041$ on top of the best decay alone.
• The chosen model barely depends on training length. Both runs pick dropout 0.5 with decay 3.0, and their test cross-entropies are 0.776 and 0.782. Without either, the test cross-entropy goes from 1.083 to 1.368 as training gets longer. Once the model is properly regularized, the step count stops acting as a hidden regularizer (see Features and Overfitting).

Why many language models pretrain without it. Dropout fights overfitting, and a model that sees each training token once doesn’t overfit in the usual sense: there’s no repeated data to memorize. PaLM’s paper says it plainly: the model “was trained without dropout, although dropout of 0.1 is used for finetuning in most cases.” Fine-tuning is the opposite situation, a small dataset seen many times, which is exactly where this post’s router lives.

Why normalize inside the network

The He initialization of Initialization and Optimizers sets every layer’s output to the right scale at the start. Training then changes the weights, and nothing keeps the scales right. A layer whose weights grow sends larger values to the next layer, whose effective learning conditions change, and so on through the depth. Learning Is Optimization showed what badly scaled inputs do to gradient descent; inside a deep network, every layer’s input is the previous layer’s output, and its scale drifts.

A normalization layer re-standardizes a set of values to mean 0 and variance 1 (or just to unit size), every forward pass, then lets the model choose the scale it wants with a learned gain. Which set of values gets standardized is the whole difference between the three kinds.

BatchNorm, and why language models don’t use it

BatchNorm (Ioffe and Szegedy, 2015) normalizes each feature across the batch. For feature $j$ over a batch of $B$ examples:

followed by a learned scale $\gamma_j$ and shift $\beta_j$. (Here $\epsilon$ is a small constant, as in Adam, and $\beta$ is a shift, not a momentum.)

The difference between the kinds of normalization is only which numbers get pooled into one mean and variance:

It worked very well for image networks, and its dependence on the batch is exactly why it doesn’t fit language models. batchnorm_flaws.py takes one He-initialized layer’s pre-activations Z on all 200 router questions and runs four tests:

# 1. The same question, normalized inside two different batches of 16.
others = rng.permutation(np.arange(1, 200))
out1 = batchnorm(Z[np.r_[0, others[:15]]], g, b)[0][0]
out2 = batchnorm(Z[np.r_[0, others[15:30]]], g, b)[0][0]
print(f"one question, two batches: its BatchNorm outputs differ by {np.abs(out1 - out2).mean():.2f} "
      f"on average (their typical size is {np.abs(out1).mean():.2f})")

# 2. How far a batch's mean is from the mean over all 200 questions, by batch size.
full_mu, full_sd = Z.mean(0), Z.std(0)
for B in (2, 8, 32, 128):
    gaps = [np.abs(Z[rng.choice(200, B, replace=False)].mean(0) - full_mu).mean() / full_sd.mean()
            for _ in range(500)]
    print(f"batch of {B:3d}: batch mean is off by {np.mean(gaps):.2f} standard deviations on average")

# 3. At inference, a batch of one: every feature equals its own mean.
print(f"batch of 1: BatchNorm output {np.abs(batchnorm(Z[:1], g, b)[0]).max():.1f} everywhere; "
      f"LayerNorm output still varies, size {np.abs(layernorm(Z[:1], g, b)[0]).mean():.2f}")

# 4. A sequence: normalize 8 token vectors together (over positions), then change
# only the *last* token. Does the *first* token's output change?
seq = rng.normal(size=(8, 64))
later = seq.copy(); later[-1] += 3.0
bn_gap = np.abs(batchnorm(seq, g, b)[0][0] - batchnorm(later, g, b)[0][0]).max()
ln_gap = np.abs(layernorm(seq, g, b)[0][0] - layernorm(later, g, b)[0][0]).max()
print(f"change the last token only: first token's output moves by {bn_gap:.3f} with BatchNorm, "
      f"{ln_gap:.3f} with LayerNorm")

Full script: batchnorm_flaws.py

one question, two batches: its BatchNorm outputs differ by 0.37 on average (their typical size is 0.91)
batch of   2: batch mean is off by 0.55 standard deviations on average
batch of   8: batch mean is off by 0.28 standard deviations on average
batch of  32: batch mean is off by 0.13 standard deviations on average
batch of 128: batch mean is off by 0.04 standard deviations on average
batch of 1: BatchNorm output 0.0 everywhere; LayerNorm output still varies, size 0.77
change the last token only: first token's output moves by 1.294 with BatchNorm, 0.000 with LayerNorm

(norms.py is below.)

• An example’s output depends on who else is in the batch. The same router question, normalized in two different batches of 16, comes out different by 0.37 on average, against a typical size of 0.91.
• Small batches give noisy statistics. With 2 examples, the batch mean is off by half a standard deviation; with 128, by 0.04. When memory forces small batches, BatchNorm’s statistics become noise.
• At inference there may be one example. A batch of one has every value equal to its own mean, so BatchNorm outputs its shift $\beta$ (here 0) everywhere and the input is erased. The standard fix keeps running averages of $\mu$ and $\sigma$ during training and uses those at test time, so the network computes a different function in training than in serving.
• Across a sequence, it leaks the future. In training, the statistics are taken over the positions of the sequences in the batch. Then changing the last token changes the first token’s output (by 1.29 here, the worst case of a batch of one sequence). A language model predicting token $t$ must not see token $t + 1$, and during training BatchNorm gives it a channel to peek. Running statistics remove the leak at inference, at the price of the train/serve mismatch above.

LayerNorm: normalize each example by itself

LayerNorm (Ba, Kiros and Hinton, 2016) turns the direction around: it normalizes each example’s features together, one row at a time. For one vector $x$ of $D$ features:

Nothing depends on other examples, so it behaves the same with a batch of one, in training and in serving, and the last line of the BatchNorm demo shows it doesn’t leak across positions: the first token’s output moved by 0.000. Worked on $x = (2, 4, 6, 8)$: $\mu = 5$, $\sigma^2 = (9 + 1 + 1 + 9)/4 = 5$, so $\hat x = (-3, -1, 1, 3)/\sqrt 5 = (-1.342, -0.447, 0.447, 1.342)$. The code, from norms.py, with its backward pass (the file also holds batchnorm, and when run prints the worked vectors and checks both backward passes against finite differences):

EPS = 1e-5

def layernorm(x, g, b):
    """Normalize each row to mean 0, variance 1, then scale by g and shift by b."""
    mu = x.mean(-1, keepdims=True)
    var = x.var(-1, keepdims=True)
    xhat = (x - mu) / np.sqrt(var + EPS)
    return xhat * g + b, (xhat, var)

def layernorm_back(dy, g, cache):
    xhat, var = cache
    dxhat = dy * g
    # The gradient of a normalization removes the parts of dxhat that would
    # only move the mean or the overall scale, which the normalization undoes.
    dx = (dxhat - dxhat.mean(-1, keepdims=True)
          - xhat * (dxhat * xhat).mean(-1, keepdims=True)) / np.sqrt(var + EPS)
    return dx, (dy * xhat).sum(0), dy.sum(0)            # dx, dg, db

def rmsnorm(x, g):
    """Divide each row by its root mean square, then scale by g. No mean, no shift."""
    rms = np.sqrt((x ** 2).mean(-1, keepdims=True) + EPS)
    xhat = x / rms
    return xhat * g, (xhat, rms)

def rmsnorm_back(dy, g, cache):
    xhat, rms = cache
    dxhat = dy * g
    dx = (dxhat - xhat * (dxhat * xhat).mean(-1, keepdims=True)) / rms
    return dx, (dy * xhat).sum(0)                        # dx, dg

Full script: norms.py

x = [2. 4. 6. 8.]:  LayerNorm [-1.342 -0.447  0.447  1.342]   RMSNorm [0.365 0.73  1.095 1.461]
x = [102. 104. 106. 108.]:  LayerNorm [-1.342 -0.447  0.447  1.342]   RMSNorm [0.971 0.99  1.009 1.028]
LayerNorm backward vs finite differences: largest relative error 1.9e-09
RMSNorm backward vs finite differences: largest relative error 2.4e-10

The backward pass has a structure worth seeing. The gradient arriving at $\hat x$, call it $\delta_{\hat x}$, has two parts the normalization will undo anyway: any push that moves every feature up together (which the mean subtraction cancels) and any push that scales the whole vector (which the division cancels). So the backward pass removes exactly those two parts, the mean of $\delta_{\hat x}$ and its component along $\hat x$, and divides by $\sigma$ (here $\sqrt{\sigma^2 + \epsilon}$, as in the code):

That’s layernorm_back, and the finite-difference check agrees to $10^{-9}$.

RMSNorm: drop the mean

RMSNorm (Zhang and Sennrich, 2019) keeps only the rescaling:

No mean, no shift $\beta$. On $(2, 4, 6, 8)$ the RMS is $\sqrt{(4 + 16 + 36 + 64)/4} = \sqrt{30} = 5.48$, giving $(0.365, 0.730, 1.095, 1.461)$. The second line of norms.py’s output shows the difference that dropping the mean makes. Shift every feature by 100 and LayerNorm’s output doesn’t change at all, while RMSNorm’s is squeezed toward $(1, 1, 1, 1)$, because the vector is now mostly its common offset. Zhang and Sennrich’s hypothesis was that the re-centering isn’t what makes normalization work, and that the rescaling alone gives the same results for less. Their experiments supported it, and LLaMA adopted RMSNorm with the words “we use the RMSNorm normalizing function”; many open language models since have followed.

What it saves. Not math:

# One normalization over a training batch of 8 x 8,192 tokens at width 4,096,
# in BF16 on an H100 SXM, against one of the layer's matmuls.
tokens, D, byte = 8 * 8192, 4096, 2
vector_flops = 67e12        # FP32 on the ordinary cores, where a norm's math runs
tensor_flops = 989e12       # BF16 on the tensor cores: matmuls only
bandwidth = 3.35e12
x_bytes = tokens * D * byte                           # one full read (or write) of the activations
for name, flops_per_elem, passes_unfused in [("LayerNorm", 8, 4), ("RMSNorm", 5, 3)]:
    t_math = tokens * D * flops_per_elem / vector_flops
    print(f"{name:9s} math {t_math * 1e6:5.0f} us   memory: unfused {passes_unfused} passes "
          f"{passes_unfused * x_bytes / bandwidth * 1e3:.2f} ms, fused 2 passes {2 * x_bytes / bandwidth * 1e3:.2f} ms")
matmul = 2 * tokens * D * D / tensor_flops
print(f"one {D} x {D} matmul on the same tokens, at tensor-core peak: {matmul * 1e3:.2f} ms")

Full script: norm_cost.py

LayerNorm math    32 us   memory: unfused 4 passes 0.64 ms, fused 2 passes 0.32 ms
RMSNorm   math    20 us   memory: unfused 3 passes 0.48 ms, fused 2 passes 0.32 ms
one 4096 x 4096 matmul on the same tokens, at tensor-core peak: 2.22 ms

A norm layer does a handful of operations per number it reads. That math runs on the GPU’s ordinary FP32 units, not the tensor cores, which only do matrix multiplications, and takes 20–32 µs. Moving the activations through memory takes 320 µs even in the best case. So a norm is memory-bound by a factor of ten to twenty, like decode in the cost sketch of AI Engineering, from the Inside. What RMSNorm saves is a pass over memory. Computed naively, LayerNorm reads the activations once for the mean, again for the variance, and again to normalize, then writes the result: four passes to RMSNorm’s three. (That’s a simplified count. A careful LayerNorm gets the mean and variance in one read, and unfused library code often makes more passes than either.) A fused kernel does either in one read and one write, and then the two cost almost the same. The backward pass saves a little more, with no mean term and no shift gradient. Zhang and Sennrich measured 7% to 64% less running time across their models; how much of that survives depends on the implementation. (Next to the layer’s matrix multiplications, 2.2 ms at peak for just one of them, even the unfused norm is small, but it isn’t free.)

What normalization does to training

Put a LayerNorm before every ReLU in the 20-layer test of Initialization and Optimizers and try the three starting scales again (deep_signal.py):

def signal(std, depth=20, width=256, seed=0):
    """The 20-layer test from "Initialization and Optimizers", with a LayerNorm before every ReLU."""
    rng = np.random.default_rng(seed)
    h = rng.normal(size=(512, width))
    Ws = [rng.normal(0, std, (width, width)) for _ in range(depth)]
    g, b = np.ones(width), np.zeros(width)
    caches, hs = [], []
    for W in Ws:
        hs.append(h)
        z, c = layernorm(h @ W, g, b)
        caches.append((z, c))
        h = np.maximum(z, 0)
    top = np.sqrt((h ** 2).mean())
    d = rng.normal(size=h.shape)
    wgrad = []
    for W, (z, c), hin in zip(Ws[::-1], caches[::-1], hs[::-1]):
        d = layernorm_back(d * (z > 0), g, c)[0]
        wgrad.append(np.sqrt(((hin.T @ d) ** 2).mean()))
        d = d @ W.T
    return top, np.sqrt((d ** 2).mean()), wgrad[-1]

print(f"{'weight std':>10s}  {'activation size, layer 20':>26s}  {'gradient reaching layer 1':>26s}  "
      f"{'gradient on W, layer 1':>23s}")
for std in (0.01, np.sqrt(2 / 256), 1.0):
    a, gx, gw = np.median([signal(std, seed=s) for s in range(5)], axis=0)
    print(f"{std:10.3f}  {a:26.2f}  {gx:26.2e}  {gw:23.2e}")

Full script: deep_signal.py

weight std   activation size, layer 20   gradient reaching layer 1   gradient on W, layer 1
     0.010                        0.71                    7.84e-01                 1.11e+02
     0.088                        0.71                    7.85e-01                 1.25e+01
     1.000                        0.71                    7.85e-01                 1.11e+00

The first two columns no longer depend on the starting scale at all. Every layer’s output is re-standardized, so the activations reach layer 20 at the same size (0.71, the size of a ReLU of a standard normal), and the gradient reaching layer 1 is the same too. The vanishing and exploding signals from Initialization and Optimizers are gone, whatever the initialization.

The last column is the catch. The network’s output no longer depends on how large $W$ is, only on its direction: scaling $W$ by 10 scales the values entering the norm by 10, and the norm divides that out. For a function that doesn’t change when $W$ is scaled, the gradient has to shrink as $W$ grows, in exact proportion: the gradient on $W$ is 111 at init std 0.01 and 1.11 at init std 1, a factor of 100 for a factor of 100 in scale. A gradient step changes $W$ by $\eta\,\|\nabla\|$, which relative to $W$‘s own size is $\eta\,\|\nabla\| / \|W\| \propto \eta / \|W\|^2$. The effective learning rate depends on the size of the weights. Small weights under a norm make every step huge.

That’s exactly what happens to the ten-layer router of Initialization and Optimizers:

grids = {"SGD": (SGD, [1e-4, 1e-3, 1e-2, 1e-1, 1.0]), "Adam": (Adam, [1e-5, 1e-4, 1e-3, 1e-2])}
out = {}
for seed in range(5):
    ((X, y), _, _), _ = four_way(*split(seed))
    sizes = [X.shape[1]] + [64] * 10 + [4]
    for start, std in [("std 0.01", lambda m: 0.01), ("He", lambda m: np.sqrt(2 / m))]:
        for norm in (None, "layer", "rms"):
            for opt_name, (make, etas) in grids.items():
                for eta in etas:
                    rng = np.random.default_rng(seed)
                    layers = init(sizes, rng, std, norm)
                    with np.errstate(all="ignore"):
                        seen = fit(layers, make(eta), X, y, 1000, rng, norm, every={1000})
                    out.setdefault((start, norm or "none", opt_name, eta), []).append(seen[1000])

Full script: norm_lr_grid.py

Each cell is the median over the 5 splits, and a run whose loss overflowed to NaN prints as inf.

start     norm   optimizer  train loss after 1,000 steps, by learning rate
std 0.01  none   SGD    0.0001: 1.386  0.001: 1.386  0.01: 1.386  0.1: 1.387  1: 1.392
std 0.01  none   Adam   1e-05: 1.386  0.0001: 1.386  0.001: 1.386  0.01: 1.387
std 0.01  layer  SGD    0.0001: 0.872  0.001: 0.069  0.01: 1.387  0.1: 1.387  1: 1.392
std 0.01  layer  Adam   1e-05: 0.808  0.0001: 0.072  0.001: 1.386  0.01: 1.387
std 0.01  rms    SGD    0.0001: 0.798  0.001: 0.062  0.01: 1.386  0.1: 1.387  1: 1.392
std 0.01  rms    Adam   1e-05: 0.722  0.0001: 0.055  0.001: 1.196  0.01: 1.387
He        none   SGD    0.0001: 1.385  0.001: 1.286  0.01: 0.002  0.1: 0.000  1: inf
He        none   Adam   1e-05: 1.230  0.0001: 0.001  0.001: 0.000  0.01: 0.000
He        layer  SGD    0.0001: 1.214  0.001: 0.226  0.01: 0.006  0.1: 0.001  1: 1.392
He        layer  Adam   1e-05: 0.803  0.0001: 0.010  0.001: 0.000  0.01: 0.506
He        rms    SGD    0.0001: 1.221  0.001: 0.207  0.01: 0.005  0.1: 0.001  1: 1.392
He        rms    Adam   1e-05: 0.810  0.0001: 0.009  0.001: 0.000  0.01: 0.000

• Without a norm, the start at init std 0.01 is hopeless. No learning rate, for either optimizer, moves it off $\log 4$.
• With a norm, it trains, at a smaller learning rate. LayerNorm or RMSNorm gets it to between 0.055 and 0.072, with SGD at $10^{-3}$ or Adam at $10^{-4}$. At the learning rates that suit the He start, it’s stuck, because every step is too large for its tiny weights. The best SGD rate is 10 to 100 times smaller than for He (the grid resolves only factors of 10, and the He start’s best two SGD rates tie), toward the squared ratio of the weight scales, about $18^2 \approx 300$. Adam’s steps are measured in the parameters’ own units (see Initialization and Optimizers), so for Adam the ratio is closer to the ratio of scales itself: 10 here.
• From a good start, the norm changes little. With He initialization, all three reach near-zero training loss at their best rates.

Why can the tiny start train at all, if its effective learning rate is so large? Because the effect corrects itself. For a layer whose output ignores the scale of $W$, the gradient is always perpendicular to $W$, so a gradient step can only make $W$ longer: $\|W_{t+1}\|^2 = \|W_t\|^2 + \eta^2\|\nabla\|^2$. Too-large steps grow the weights until the steps are no longer too large, a mechanism Arora, Li and Lyu (2019) called auto rate-tuning. It works here only once the learning rate is small enough for that to happen before training goes wrong.

The same scale invariance changes what weight decay does in a normalized network (van Laarhoven, 2017). Shrinking $W$ doesn’t change the function the layer computes, so decay no longer penalizes complexity directly. What it does is hold back the growth just described, which keeps the effective learning rate from shrinking; decay and growth settle into a balance. It’s one more reason Initialization and Optimizers found that $\lambda$, $\eta$ and training length act together. The norms’ own gains are different: they aren’t scale-invariant, since shrinking $\gamma$ shrinks the output directly, so they, like biases, are usually exempt from decay, as optim.decays does.

Where the norms go in a Transformer

In a Transformer the norm sits at a fixed place in every block. The Residual Stream describes the layout. The original 2017 Transformer normalized after adding each sublayer’s output to the stream (post-norm). By 2018–19, implementations and papers (Wang et al., 2019) had moved it to the input of each sublayer (pre-norm), so the stream itself is never normalized. GPT-2 adopted it and added one final norm before the unembedding, needed because the un-normalized stream keeps growing through the layers. Xiong et al. (2020) showed why that matters for optimization: with post-norm, the gradients near the output are large at initialization, and training needs a careful warmup (see Initialization and Optimizers) to survive the first steps. Pre-norm is much less sensitive. Most modern open language models use pre-norm with RMSNorm, some with extra norms, on a sublayer’s output or on the attention queries and keys. Build a Tiny GPT, End to End, and Train It on This Blog builds it.

Where this leads

• PyTorch from the Inside replaces these hand-written backward passes with PyTorch’s autograd, and nn.LayerNorm, nn.RMSNorm and nn.Dropout with their library versions, including the fused kernels behind the cost numbers.
• The Residual Stream puts pre-norm RMSNorm and residual connections into a Transformer block, where they let dozens of layers train stably.
• Build a Tiny GPT, End to End, and Train It on This Blog trains such blocks on this blog’s small corpus, with dropout on the attention weights, and chooses the dropout on a dev slice.

Takeaways

• Inverted dropout keeps each unit with probability $1 - p$ and divides by $1 - p$, so the average signal is unchanged and nothing changes at test time. It trains an implicit crowd of thinner networks.
• Dropout needs a limit on the logits. Without weight decay it widened the router’s margins, made it more confident when wrong, and worsened cross-entropy. With decay, and training long enough to overfit, it improved test cross-entropy by 0.19 beyond the best decay alone. Accuracy changed within its error throughout.
• Normalization is the feature scaling of Learning Is Optimization inside the network, every step. BatchNorm normalizes each feature across the batch, which ties examples together, fails with a batch of one, and leaks future tokens across a sequence. LayerNorm normalizes each example across its features and has none of those problems.
• RMSNorm drops the mean and the shift. Unfused, that saves a pass over memory; fused, the two cost about the same, because norms are bound by memory traffic (by 10–20×), not math.
• LayerNorm’s backward pass removes the parts of the gradient that the normalization would undo: its mean and its component along $\hat x$. RMSNorm’s removes only the component along $\hat x$.
• Under a norm, the network’s output ignores the scale of $W$, the gradient on $W$ shrinks as $1/\|W\|$, and the effective learning rate scales as $\eta/\|W\|^2$. A norm rescued the router’s tiny start, but only at a 10–100 times smaller learning rate (Adam and SGD).
• Language models use pre-norm RMSNorm in every block; many pretrain with no dropout, and use it again for fine-tuning.

References

Start with Ba, Kiros and Hinton for LayerNorm and Xiong et al. for why pre-norm matters. Srivastava et al. (2014) is the full dropout paper, with the ensemble view.

• Hinton, Srivastava, Krizhevsky, Sutskever and Salakhutdinov, Improving Neural Networks by Preventing Co-Adaptation of Feature Detectors (2012): the original dropout paper, which also dropped inputs.
• Srivastava, Hinton, Krizhevsky, Sutskever and Salakhutdinov, Dropout: A Simple Way to Prevent Neural Networks from Overfitting (JMLR, 2014): dropout as an ensemble, and the weight-scaling rule.
• Ioffe and Szegedy, Batch Normalization (ICML, 2015): BatchNorm and its running statistics.
• Ba, Kiros and Hinton, Layer Normalization (2016): the LayerNorm definition, normalizing each example over its features.
• Zhang and Sennrich, Root Mean Square Layer Normalization (NeurIPS, 2019): RMSNorm, and its 7%–64% running-time savings.
• Arora, Li and Lyu, Theoretical Analysis of Auto Rate-Tuning by Batch Normalization (ICLR, 2019): why scale-invariant layers correct their own step size.
• van Laarhoven, L2 Regularization versus Batch and Weight Normalization (2017): scale invariance, and weight decay as an effective-learning-rate control.
• Xiong et al., On Layer Normalization in the Transformer Architecture (ICML, 2020): pre-norm vs post-norm, and warmup.
• Wang et al., Learning Deep Transformer Models for Machine Translation (ACL, 2019): pre-norm for deep Transformers.
• Radford et al. (OpenAI), Language Models are Unsupervised Multitask Learners (2019): GPT-2, with pre-norm and a final norm.
• Touvron et al. (Meta AI), LLaMA: Open and Efficient Foundation Language Models (2023): the source for LLaMA’s pre-normalization with RMSNorm.
• Chowdhery et al. (Google), PaLM: Scaling Language Modeling with Pathways (2022): pretraining without dropout, and 0.1 for fine-tuning.