Layer normalization over the last axis
Easy9 min11 hidden tests
Every block of a chat model's transformer starts with layer normalization: each token's vector of d features is rescaled to mean 0 and variance 1, then a learned gain gamma and bias beta set its scale and offset. Write layer_norm(x, gamma, beta, eps) that does this over the last axis of x, whatever its other dimensions are. The token [1.0, 2.0, 3.0] with gain 1 and bias 0 becomes about [-1.2247, 0.0, 1.2247].
Normalize one token
y = (x − mean) / √(var + eps) × gamma + betaThis token’s 6 features have mean 1.5 and variance 2.1, the mean squared gap divided by d = 6.
Stage 1 of 3.
Examples
Mean 2 and variance 2/3
layer_norm([[1.0, 2.0, 3.0]], [1.0, 1.0, 1.0], [0.0, 0.0, 0.0])
returns [[-1.22474, 0, 1.22474]]
A constant token comes out as beta
layer_norm([[2.0, 2.0], [0.0, 4.0]], [2.0, 1.0], [0.5, -0.5])
returns [[0.5, -0.5], [-1.5, 0.499999]]
Requirements
xhas shape (..., d): any number of leading axes, then d features.gammaandbetahave shape (d,).- For each vector along the last axis,
y = (x - mean) / sqrt(var + eps) * gamma + beta. varis the population variance: the mean of the squared deviations, divided by d (NumPy'snp.vardefault).- Return a float array with the same shape as
x. A vector whose features are all equal comes out asbeta.
The idea behind it
Layer normalization (LayerNorm) rescales each example's feature vector to mean 0 and variance 1, then applies a learned per-feature gain and bias: y = (x - mean) / sqrt(var + eps) * gamma + beta. The statistics come from one example's own features, so the result does not depend on the batch. It behaves the same with a batch of 1, during generation and in training. That is why transformers use it instead of batch normalization. The eps, often 1e-5, keeps the division safe when all features are equal. The variance is the population one, dividing by d. Keeping each layer's inputs at a steady scale lets deep stacks train with larger learning rates. Many recent LLMs use RMSNorm, a cousin that skips the mean and divides by the root mean square.