Layer Normalization in deep learning

Module 02 · Deep Learning · 25 concepts

In depth / 15 lines

Layer Normalization, in more detail

Layer normalization does the same rescaling as batch normalization, along a different axis.

It normalises across the features of one example, using only that example.

No other row in the batch is involved at any point.

That single difference is what makes it suit sequences and small batches.

It works identically with a batch size of one, which batch normalization cannot manage.

It behaves the same in training and inference, so there are no running averages to drift.

Sequences of different lengths cause it no trouble, since positions are handled independently.

Like batch normalization it learns a scale and a shift to restore whatever it removed.

It is applied at every block of a transformer, usually twice.

Pre-norm places it before the attention and feed-forward parts, which trains more stably.

Post-norm, in the original paper, places it after and needs careful warmup to survive.

RMSNorm is a lighter variant that divides by the root mean square and skips the mean.

It is slightly cheaper and works about as well, so large recent models tend to use it.

Choose layer normalization for sequence models and batch normalization for convolutional vision models.

Both exist to keep the numbers flowing through a deep network within a sane range.

Want the full technical article? Wikipedia ↗

Deep dive

What is underneath Layer Normalization

The same operation, over a different axis

Layer normalisation standardises across the features of a single example, rather than across the examples of a single feature.

\mu = \frac{1}{d} \sum_{j=1}^{d} x_j \qquad \sigma^2 = \frac{1}{d} \sum_{j=1}^{d} (x_j - \mu)^2

y = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta

$d$ : the width of the layer, the number of features in one vector.
$\gamma, \beta$ : learned scale and shift, still one pair per feature.
$\odot$ : elementwise multiplication.

The formula is nearly identical to batch normalisation. The index it sums over is the entire difference, and it changes everything.

Why that axis wins for sequences

The statistics come from one example only, so the batch size is irrelevant.

That means training and inference compute exactly the same function, with no running averages to keep and nothing to converge.

It also means a batch of one behaves identically to a batch of a thousand, which is what inference actually looks like.

And it survives variable-length sequences: a padded position does not contaminate a real one, because no statistic is shared between positions.

Those four properties are why every transformer uses this and not batch normalisation.

RMSNorm, which most recent models use instead

y = \gamma \odot \frac{x}{\sqrt{\frac{1}{d}\sum_{j=1}^{d} x_j^2 + \epsilon}}

The mean subtraction is dropped entirely, leaving only the rescaling by root-mean-square.

It works about as well, has no $\beta$ , and needs one pass over the vector instead of two. At the scale transformers run at, that is a measurable saving.

One forward pass, in order

Take one position's vector of $d$ features.
Compute its mean and variance over those features.
Subtract the mean, divide by the root of the variance plus epsilon.
Scale by gamma and add beta.
Move to the next position, sharing nothing with this one.

What it actually does for training

The usual story is that it reduces internal covariate shift, and the evidence for that is weak.

The better-supported account is that it smooths the loss surface, which makes the gradient more predictable from step to step and lets a larger learning rate be used safely.

Test yourself

One question on this card. Every option explains itself once you pick, right or wrong.

Question 1

Why is layer normalization used in transformers instead of batch normalization?

Deep Learning

25 concepts in this module, then a graded test with every answer explained.

Study this module →