Neural Networks & LLM Components · Medium

Layer Normalization

Implement LayerNorm and compare its Pre-LN and Post-LN placement in a residual block.

Task

Implement layer_normalization(values: list[list[float]], gamma: list[float], beta: list[float], sublayer_scale: list[float], placement: str, epsilon: float = 1e-5) -> list[list[float]]. Implement a toy residual block with F(z) = z * sublayer_scale. For "pre", return x + F(LN(x)); for "post", return LN(x + F(x)).

Requirements

Every row, gamma, beta, and sublayer_scale have the same feature width.
Compute a separate mean and population variance for each row.
Apply epsilon inside the square root before gamma and beta.
When placement is "pre", normalize before F and add the result to the original x.
When placement is "post", apply F to x, add the residual, and normalize the sum.
Preserve empty batches and zero-width rows without producing NaN.

Example

values = [[1.0, 3.0]]
gamma = [1.0, 1.0]
beta = [0.0, 0.0]
sublayer_scale = [2.0, 0.5]

layer_normalization(values, gamma, beta, sublayer_scale, "pre", epsilon=0.0)
# [[-1.0, 3.5]]

layer_normalization(values, gamma, beta, sublayer_scale, "post", epsilon=0.0)
# [[-1.0, 1.0]]