DL0034 Layer Norm

What is layer normalization, and why is it used in Transformers?

Answer

Layer Normalization (LN) standardizes the features of each individual sample: for one token’s embedding vector, it computes the mean and variance across the feature dimension only, rescales to zero mean and unit variance, then applies a learnable scale and shift. Unlike BatchNorm, it never looks across the batch, which is exactly why Transformers, with variable-length sequences and small or on-the-fly batches, rely on it in every block.

(1) Normalization Within a Sample: Mean and variance come from the d_{model} features of a single token: one set of statistics per token, not per batch.
(2) Batch-Size Independence: Behavior is identical at train and test time and for any batch size: no running statistics, no mismatch.
(3) Stabilizes Training: Keeps activations in a consistent range, preventing exploding/vanishing gradients and enabling deep stacks to converge faster.

3D illustration of Layer Normalization: each sample (row) is normalized across its own feature dimensions, independent of the other samples in the batch.

Figure 1: Layer Normalization standardizes each sample (row) across its own feature dimensions; statistics never cross sample boundaries, so batch size is irrelevant.

Mathematical Formulation:
\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta
\mu = \frac{1}{d}\sum_{i=1}^{d} x_i
\sigma^2 = \frac{1}{d}\sum_{i=1}^{d}(x_i - \mu)^2

Where:

  • x_i is one feature of a single token’s d-dimensional vector; statistics are computed over i = 1..d for that token alone.
  • \epsilon is a small constant for numerical stability; \gamma, \beta are learnable per-feature scale and shift that restore representational freedom.

Why Not BatchNorm: BatchNorm’s statistics mix information across samples, degrade with small or variable-size batches, behave differently at train vs inference, and pad tokens corrupt the per-feature means of variable-length sequences: all fatal for typical Transformer workloads.

Two mini flowcharts of a Transformer sub-layer contrasting Post-LN where normalization follows the residual addition with Pre-LN where normalization precedes the sub-layer and the residual path stays clean.

Figure 2: Placement variants: original Post-LN (after the residual add) vs modern Pre-LN (inside the residual branch), which keeps a clean gradient highway.

Where It Sits: Every attention and FFN sub-layer is wrapped as \mathrm{LN}(x + \mathrm{Sublayer}(x)) (Post-LN) in the original paper; most modern LLMs use Pre-LN, normalizing the sub-layer input instead, which trains stably even without learning-rate warmup.


Login to view more content


Log in to track your progress

Comments

Leave a Reply

Your email address will not be published. Required fields are marked *