Layer Normalization

Your program will rescale one token's vector so its values have mean 0 and spread 1, then let learned numbers (gamma, beta) decide how much of that scale to keep. You will also write RMSNorm, the cheaper variant Llama uses.

Why a transformer needs it

Every block adds something to the running vector, and the next block multiplies it by weight matrices. Over dozens of blocks the numbers drift: some coordinates grow into the hundreds, others shrink toward zero, and a weight matrix tuned for one scale is wrong for another. Normalising the input of each sublayer pins the scale, so each sublayer sees data with the same statistics no matter how deep it sits.

The recipe, for one token

The input is one vector x of D numbers (one row of the matrix from lesson 2). Normalise across its D features, one token at a time:

mean = (x[0] + ... + x[D-1]) / D
var  = ((x[0]-mean)^2 + ... + (x[D-1]-mean)^2) / D
y[i] = gamma[i] * (x[i] - mean) / sqrt(var + eps) + beta[i]

eps (1e-5 in GPT-2) sits inside the square root. It keeps the division finite when every coordinate is equal and var is 0. gamma and beta are two learned vectors of length D: gamma rescales each coordinate, beta shifts it, so the model can undo the normalisation wherever that helps. A fresh model starts with gamma = 1, beta = 0.

Worked trace, x = [2, 4, 6, 8], gamma = [2, 1, 1, 0.5], beta = [0, 0, 1, 0], eps = 1e-5:

mean = 20 / 4 = 5
deviations      = -3, -1, 1, 3
var  = (9 + 1 + 1 + 9) / 4 = 5          sqrt(5.00001) = 2.2361
normalised      = -1.3416, -0.4472, 0.4472, 1.3416
y = [2*-1.3416, 1*-0.4472, 1*0.4472 + 1, 0.5*1.3416]
  = [-2.6833, -0.4472, 1.4472, 0.6708]

Without gamma and beta the output has mean 0 and variance 1 for any input. A constant vector like [7, 7, 7] normalises to all zeros, and the output is then exactly beta.

LayerNorm is not BatchNorm

BatchNorm computes a mean and variance for each feature across the examples in a batch. That makes the result depend on who else is in the batch, which breaks at batch size 1 and on sequences of different lengths. LayerNorm uses only the one token's own D numbers, so training and generation behave identically.

Where it goes

The original transformer put the norm after each residual addition (post-norm). GPT-2 and almost everything since put it before each sublayer (pre-norm): x = x + sublayer(norm(x)). Lesson 9 builds both and shows why pre-norm survives deep stacks.

RMSNorm

Llama drops the centring and the shift and divides by the root mean square alone:

y[i] = gamma[i] * x[i] / sqrt( (x[0]^2 + ... + x[D-1]^2) / D + eps )

For x = [2, -2, 4, -4] the mean of squares is 10, so y = [0.6325, -0.6325, 1.2649, -1.2649] * gamma. It saves a pass over the data and loses little in practice.

Common mistakes

• Putting eps outside the root, or adding it to the standard deviation instead of the variance. Both agree for tiny eps and disagree as soon as eps is large.
• Dividing the variance by D - 1. Some libraries default to that unbiased estimate; LayerNorm uses D.
• Computing the variance as mean(x^2) - mean^2. With values near 1e8 the two terms are almost equal and the subtraction wipes out the answer. Subtract the mean first, then square.
• Applying RMSNorm's formula after subtracting the mean, or LayerNorm's without it.

Your exercise

Read EPS, GAMMA and BETA lines, then answer every LN line with LayerNorm and every RMS line with RMSNorm, one output line each. The starter does the parsing and printing; you write the two functions.