AI Engineer Exhaustive
EasyPlusGPT-2

Layer Normalization

GPT-2

Easy

Pre-norm layer normalization

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Layer normalization (Ba, Kiros & Hinton, 2016) is the only normalization in GPT-2. Section 2.3 says two things: LN is moved to the input of each sub-block, “similar to a pre-activation residual network,” and an extra LN is placed after the final self-attention block. That is pre-norm plus a terminal ln_f, not the post-norm GPT-1 / original Transformer pattern where you add first and normalize the sum.

Pre-norm leaves the residual path as an identity. Gradients can travel down (H^{(0)}+\sum_{\ell}(A^{(\ell)}+M^{(\ell)})) without passing through an LN Jacobian at every add. That is the paper’s stated reason for the move, together with the residual init scale, when depth goes to 48 layers.

How it works

The official norm reduces over the last axis (the channel / n_embd axis), not over batch or time. For a slice (x\in\mathbb{R}^{d}),

\mu=\frac{1}{d}\sum_{j=1}^{d}x_j,\qquad s=\frac{1}{d}\sum_{j=1}^{d}(x_j-\mu)^{2},\qquad \hat x=(x-\mu)\,(s+\varepsilon)^{-1/2},

\mathrm{LN}(x)=g\odot\hat x+b.

g and b are vectors of length n_state (d for stream norms). They start at ones and zeros. (\varepsilon=10^{-5}). The code computes (s) as the mean of squared deviations (tf.reduce_mean(tf.square(x-u))) and rescales with tf.rsqrt(s + epsilon). That is the population (biased) variance, not Bessel-corrected sample variance, and it is a true variance: it is not RMSNorm, which would skip (\mu) and normalize by (\mathrm{mean}(x^{2})).

Tensors entering norm are [B, T, d]. Broadcasted (\mu,s) have shape [B, T, 1]. Each token is normalized independently. There is no running mean: this is not batch norm.

Placement in block:

A=\mathrm{Attn}(\mathrm{LN}_1(H)),\quad H\leftarrow H+A,

M=\mathrm{MLP}(\mathrm{LN}_2(H)),\quad H\leftarrow H+M.

ln_1 and ln_2 are distinct parameter pairs per layer (h{ℓ}/ln_1/{g,b}, h{ℓ}/ln_2/{g,b}). The residual add uses the unnormalized (H). After layer (L-1), model() applies norm(h, 'ln_f') before the tied embedding projection. Embeddings themselves are not normalized on the way in.

Because LN is inside each sub-block, the attention and MLP always see zero-mean, unit-scale channels (up to (g,b)), even if the residual stream’s magnitude grows with depth. The stream itself can drift; only the branch inputs are standardized.

A common rewrite uses (\sigma=\sqrt{s+\varepsilon}) in the denominator. That matches rsqrt as long as (\varepsilon) is added inside the square root. Adding (\varepsilon) outside, or using (10^{-6}) from another codebase, will not match official numerics exactly, though it is usually close in float32.

Official code

norm in src/model.py of openai/gpt-2. Call sites: block (ln_1 before attn, ln_2 before mlp) and model (ln_f after the layer loop). Variable names are g and b, not gamma / beta.

Watch-outs

Sources