Layer Normalization

Layer normalization rewrites one token's vector using nothing but that vector: subtract the vector's own mean, divide by its own standard deviation, then multiply by a learned gain and add a learned bias, one of each per dimension.

What it computes

$x$ is one token's row, d components wide. $\mu$ and $\sigma^2$ are that row's own mean and variance — the variance divided by d rather than by d − 1, which is unbiased=False in PyTorch and what GPT-2 was trained with. The fraction is the row standardized: mean 0 and variance 1, whatever arrived. $\gamma$ and $\beta$ are the learned gain and bias, applied componentwise.

Three details in the arithmetic are worth naming:

• The statistics are the token's own, so a token's output never depends on what else is in the batch or the sequence, and training and inference compute the same thing. This is the contrast with batch normalization, which normalizes across the batch and is not used in transformers.

• eps guards the division. A token vector whose components are all equal has variance 0, and without eps the division would be by zero. GPT-2 uses $\epsilon = 10^{-5}$, PyTorch's default.

• gain and bias are the only trainable parameters, d of each: one pair per dimension, shared by every token in the sequence and every sequence in the batch. They start at 1 and 0, where the layer is standardization and nothing more. Training moves them, so the output rows stop having mean 0 and variance 1. That is intended: standardizing fixes the scale and the offset of what leaves, and the pair then lets training pick a scale and an offset per dimension rather than inheriting whatever the input carried.

Why it matters

Switching between X_small and X_big changes the layer's input magnitude and offset; the gain and bias from the section above stay at 1 and 0 here.

The two branches are the same rows forty times apart in magnitude, each with its own offset, and layer normalization gives the same output for both — up to the eps under the square root, which does not scale with the rows and whose effect the readout prints. However large or small the values arriving have become, the next layer sees rows at the same scale, so its gradients neither explode nor vanish with them.

GPT-2 places layer normalization in three positions: before the attention sublayer and before the feed-forward network of every transformer block, and once after the last block. With 12 blocks, GPT-2 (124M) has 25 of them. Normalizing only each sublayer's input is called pre-norm: the sublayer's shortcut connection carries the input around it unnormalized. The original transformer normalized after each shortcut connection instead. With pre-norm, the output of the last block's second shortcut connection is not normalized, so the final layer normalization normalizes it before the output layer.

References

Sebastian Raschka, Build a Large Language Model (From Scratch) (Manning, 2024). Section 4.2 implements the layer in PyTorch, which the code exported above follows. It names the two parameters scale and shift; this chapter uses Ba et al.'s gain and bias instead.

Jimmy Lei Ba, Jamie Ryan Kiros and Geoffrey E. Hinton, Layer Normalization (2016). The paper the method comes from.

Alec Radford et al., Language Models are Unsupervised Multitask Learners (2019). Section 2.3 moves the normalization to the input of each sublayer and adds the final one.