Layer Normalization

Layer normalization takes the same re-centre-and-rescale idea as batch normalization but computes the statistics across the features of one sample, so it never looks at the rest of the batch. After this lesson your program normalizes every row of a batch independently, the way each transformer layer does.

Why move the axis

Batch norm needs a decent batch to estimate a mean and variance, behaves differently in training and in evaluation, and is awkward for sequences whose lengths vary or whose batch size is 1 while generating text. Layer norm sidesteps all of that: each sample is normalized using only its own D features, so training and inference run identical code and the result for a sample cannot depend on its neighbours.

The formula

For one sample x with D features:

mean   = (x1 + ... + xD) / D
var    = ((x1 - mean)^2 + ... + (xD - mean)^2) / D
x_hat  = (x - mean) / sqrt(var + eps)
y_i    = gamma_i * x_hat_i + beta_i

It looks exactly like batch norm; only the axis differs:

gamma and beta are vectors of length D here, one scale and shift per feature, so the network can still express "this feature should be larger than that one" after every feature has been put on the same footing. By default gamma is all ones and beta all zeros.

A worked trace

One sample with three features, x = [10, 20, 60]:

mean = 30
var  = (400 + 100 + 900) / 3 = 466.67        std = 21.60
x_hat = [-0.9258, -0.4629, 1.3887]

Multiply the sample by 10 to get [100, 200, 600] and the output is the same, to four decimals, because both the mean and the standard deviation scale with it. Add 1000 to every feature and again the output does not move. Layer norm discards a sample's overall offset and scale and keeps only its shape. That is exactly the stabilising effect it is used for: whatever happened upstream, the next layer sees a steady distribution.

Common mistakes

• Computing the statistics down the batch (across rows) instead of along each row.
• Dividing the variance by D - 1.
• Putting eps outside the square root, or forgetting it for constant rows. A row of equal values has variance 0 and must still produce zeros, not nan.
• Using one scalar gamma for all features when the exercise passes a vector.
• Letting one row's values affect another row's output.

Where it lives

Layer norm is the standard normalization inside transformers. The original architecture applied it after each residual addition; modern ones apply it before each sublayer (pre-norm), which you will build in the last lesson. A cheaper variant, RMSNorm, skips the mean subtraction and divides by the root-mean-square only.

Your exercise

INPUT lines add samples; NORM prints each pending sample normalized on its own, then clears them. GAMMA, BETA and EPS stay set. The starter fills in the default gamma and beta. Implement layer_norm for a single sample.